dataset.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376
  1. import os
  2. from abc import ABC, abstractmethod
  3. from typing import Callable, Iterator
  4. import fitz
  5. from loguru import logger
  6. from magic_pdf.config.enums import SupportedPdfParseMethod
  7. from magic_pdf.data.schemas import PageInfo
  8. from magic_pdf.data.utils import fitz_doc_to_image
  9. from magic_pdf.filter import classify
  10. from magic_pdf.model.sub_modules.language_detection.utils import auto_detect_lang
  11. class PageableData(ABC):
  12. @abstractmethod
  13. def get_image(self) -> dict:
  14. """Transform data to image."""
  15. pass
  16. @abstractmethod
  17. def get_doc(self) -> fitz.Page:
  18. """Get the pymudoc page."""
  19. pass
  20. @abstractmethod
  21. def get_page_info(self) -> PageInfo:
  22. """Get the page info of the page.
  23. Returns:
  24. PageInfo: the page info of this page
  25. """
  26. pass
  27. @abstractmethod
  28. def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay):
  29. """draw rectangle.
  30. Args:
  31. rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1]
  32. color (list[float] | None): three element tuple which describe the RGB of the board line, None means no board line
  33. fill (list[float] | None): fill the board with RGB, None means will not fill with color
  34. fill_opacity (float): opacity of the fill, range from [0, 1]
  35. width (float): the width of board
  36. overlay (bool): fill the color in foreground or background. True means fill in background.
  37. """
  38. pass
  39. @abstractmethod
  40. def insert_text(self, coord, content, fontsize, color):
  41. """insert text.
  42. Args:
  43. coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1]
  44. content (str): the text content
  45. fontsize (int): font size of the text
  46. color (list[float] | None): three element tuple which describe the RGB of the board line, None will use the default font color!
  47. """
  48. pass
  49. class Dataset(ABC):
  50. @abstractmethod
  51. def __len__(self) -> int:
  52. """The length of the dataset."""
  53. pass
  54. @abstractmethod
  55. def __iter__(self) -> Iterator[PageableData]:
  56. """Yield the page data."""
  57. pass
  58. @abstractmethod
  59. def supported_methods(self) -> list[SupportedPdfParseMethod]:
  60. """The methods that this dataset support.
  61. Returns:
  62. list[SupportedPdfParseMethod]: The supported methods, Valid methods are: OCR, TXT
  63. """
  64. pass
  65. @abstractmethod
  66. def data_bits(self) -> bytes:
  67. """The bits used to create this dataset."""
  68. pass
  69. @abstractmethod
  70. def get_page(self, page_id: int) -> PageableData:
  71. """Get the page indexed by page_id.
  72. Args:
  73. page_id (int): the index of the page
  74. Returns:
  75. PageableData: the page doc object
  76. """
  77. pass
  78. @abstractmethod
  79. def dump_to_file(self, file_path: str):
  80. """Dump the file
  81. Args:
  82. file_path (str): the file path
  83. """
  84. pass
  85. @abstractmethod
  86. def apply(self, proc: Callable, *args, **kwargs):
  87. """Apply callable method which.
  88. Args:
  89. proc (Callable): invoke proc as follows:
  90. proc(self, *args, **kwargs)
  91. Returns:
  92. Any: return the result generated by proc
  93. """
  94. pass
  95. @abstractmethod
  96. def classify(self) -> SupportedPdfParseMethod:
  97. """classify the dataset
  98. Returns:
  99. SupportedPdfParseMethod: _description_
  100. """
  101. pass
  102. @abstractmethod
  103. def clone(self):
  104. """clone this dataset
  105. """
  106. pass
  107. class PymuDocDataset(Dataset):
  108. def __init__(self, bits: bytes, lang=None):
  109. """Initialize the dataset, which wraps the pymudoc documents.
  110. Args:
  111. bits (bytes): the bytes of the pdf
  112. """
  113. self._raw_fitz = fitz.open('pdf', bits)
  114. self._records = [Doc(v) for v in self._raw_fitz]
  115. self._data_bits = bits
  116. self._raw_data = bits
  117. if lang == '':
  118. self._lang = None
  119. elif lang == 'auto':
  120. self._lang = auto_detect_lang(bits)
  121. logger.info(f"lang: {lang}, detect_lang: {self._lang}")
  122. else:
  123. self._lang = lang
  124. def __len__(self) -> int:
  125. """The page number of the pdf."""
  126. return len(self._records)
  127. def __iter__(self) -> Iterator[PageableData]:
  128. """Yield the page doc object."""
  129. return iter(self._records)
  130. def supported_methods(self) -> list[SupportedPdfParseMethod]:
  131. """The method supported by this dataset.
  132. Returns:
  133. list[SupportedPdfParseMethod]: the supported methods
  134. """
  135. return [SupportedPdfParseMethod.OCR, SupportedPdfParseMethod.TXT]
  136. def data_bits(self) -> bytes:
  137. """The pdf bits used to create this dataset."""
  138. return self._data_bits
  139. def get_page(self, page_id: int) -> PageableData:
  140. """The page doc object.
  141. Args:
  142. page_id (int): the page doc index
  143. Returns:
  144. PageableData: the page doc object
  145. """
  146. return self._records[page_id]
  147. def dump_to_file(self, file_path: str):
  148. """Dump the file
  149. Args:
  150. file_path (str): the file path
  151. """
  152. dir_name = os.path.dirname(file_path)
  153. if dir_name not in ('', '.', '..'):
  154. os.makedirs(dir_name, exist_ok=True)
  155. self._raw_fitz.save(file_path)
  156. def apply(self, proc: Callable, *args, **kwargs):
  157. """Apply callable method which.
  158. Args:
  159. proc (Callable): invoke proc as follows:
  160. proc(dataset, *args, **kwargs)
  161. Returns:
  162. Any: return the result generated by proc
  163. """
  164. if 'lang' in kwargs and self._lang is not None:
  165. kwargs['lang'] = self._lang
  166. return proc(self, *args, **kwargs)
  167. def classify(self) -> SupportedPdfParseMethod:
  168. """classify the dataset
  169. Returns:
  170. SupportedPdfParseMethod: _description_
  171. """
  172. return classify(self._data_bits)
  173. def clone(self):
  174. """clone this dataset
  175. """
  176. return PymuDocDataset(self._raw_data)
  177. class ImageDataset(Dataset):
  178. def __init__(self, bits: bytes):
  179. """Initialize the dataset, which wraps the pymudoc documents.
  180. Args:
  181. bits (bytes): the bytes of the photo which will be converted to pdf first. then converted to pymudoc.
  182. """
  183. pdf_bytes = fitz.open(stream=bits).convert_to_pdf()
  184. self._raw_fitz = fitz.open('pdf', pdf_bytes)
  185. self._records = [Doc(v) for v in self._raw_fitz]
  186. self._raw_data = bits
  187. self._data_bits = pdf_bytes
  188. def __len__(self) -> int:
  189. """The length of the dataset."""
  190. return len(self._records)
  191. def __iter__(self) -> Iterator[PageableData]:
  192. """Yield the page object."""
  193. return iter(self._records)
  194. def supported_methods(self):
  195. """The method supported by this dataset.
  196. Returns:
  197. list[SupportedPdfParseMethod]: the supported methods
  198. """
  199. return [SupportedPdfParseMethod.OCR]
  200. def data_bits(self) -> bytes:
  201. """The pdf bits used to create this dataset."""
  202. return self._data_bits
  203. def get_page(self, page_id: int) -> PageableData:
  204. """The page doc object.
  205. Args:
  206. page_id (int): the page doc index
  207. Returns:
  208. PageableData: the page doc object
  209. """
  210. return self._records[page_id]
  211. def dump_to_file(self, file_path: str):
  212. """Dump the file
  213. Args:
  214. file_path (str): the file path
  215. """
  216. dir_name = os.path.dirname(file_path)
  217. if dir_name not in ('', '.', '..'):
  218. os.makedirs(dir_name, exist_ok=True)
  219. self._raw_fitz.save(file_path)
  220. def apply(self, proc: Callable, *args, **kwargs):
  221. """Apply callable method which.
  222. Args:
  223. proc (Callable): invoke proc as follows:
  224. proc(dataset, *args, **kwargs)
  225. Returns:
  226. Any: return the result generated by proc
  227. """
  228. return proc(self, *args, **kwargs)
  229. def classify(self) -> SupportedPdfParseMethod:
  230. """classify the dataset
  231. Returns:
  232. SupportedPdfParseMethod: _description_
  233. """
  234. return SupportedPdfParseMethod.OCR
  235. def clone(self):
  236. """clone this dataset
  237. """
  238. return ImageDataset(self._raw_data)
  239. class Doc(PageableData):
  240. """Initialized with pymudoc object."""
  241. def __init__(self, doc: fitz.Page):
  242. self._doc = doc
  243. def get_image(self):
  244. """Return the image info.
  245. Returns:
  246. dict: {
  247. img: np.ndarray,
  248. width: int,
  249. height: int
  250. }
  251. """
  252. return fitz_doc_to_image(self._doc)
  253. def get_doc(self) -> fitz.Page:
  254. """Get the pymudoc object.
  255. Returns:
  256. fitz.Page: the pymudoc object
  257. """
  258. return self._doc
  259. def get_page_info(self) -> PageInfo:
  260. """Get the page info of the page.
  261. Returns:
  262. PageInfo: the page info of this page
  263. """
  264. page_w = self._doc.rect.width
  265. page_h = self._doc.rect.height
  266. return PageInfo(w=page_w, h=page_h)
  267. def __getattr__(self, name):
  268. if hasattr(self._doc, name):
  269. return getattr(self._doc, name)
  270. def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay):
  271. """draw rectangle.
  272. Args:
  273. rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1]
  274. color (list[float] | None): three element tuple which describe the RGB of the board line, None means no board line
  275. fill (list[float] | None): fill the board with RGB, None means will not fill with color
  276. fill_opacity (float): opacity of the fill, range from [0, 1]
  277. width (float): the width of board
  278. overlay (bool): fill the color in foreground or background. True means fill in background.
  279. """
  280. self._doc.draw_rect(
  281. rect_coords,
  282. color=color,
  283. fill=fill,
  284. fill_opacity=fill_opacity,
  285. width=width,
  286. overlay=overlay,
  287. )
  288. def insert_text(self, coord, content, fontsize, color):
  289. """insert text.
  290. Args:
  291. coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1]
  292. content (str): the text content
  293. fontsize (int): font size of the text
  294. color (list[float] | None): three element tuple which describe the RGB of the board line, None will use the default font color!
  295. """
  296. self._doc.insert_text(coord, content, fontsize=fontsize, color=color)