mixin.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Union, Tuple, List, Dict, Any, Iterator
  15. from abc import abstractmethod
  16. from pathlib import Path
  17. import mimetypes
  18. import json
  19. import copy
  20. import re
  21. import numpy as np
  22. from PIL import Image
  23. import pandas as pd
  24. from ....utils import logging
  25. from ...utils.io import (
  26. JsonWriter,
  27. ImageReader,
  28. ImageWriter,
  29. CSVWriter,
  30. HtmlWriter,
  31. XlsxWriter,
  32. TextWriter,
  33. VideoWriter,
  34. MarkdownWriter,
  35. )
  36. class StrMixin:
  37. """Mixin class for adding string conversion capabilities."""
  38. @property
  39. def str(self) -> Dict[str, str]:
  40. """Property to get the string representation of the result.
  41. Returns:
  42. Dict[str, str]: The string representation of the result.
  43. """
  44. return self._to_str()
  45. def _to_str(
  46. self,
  47. ):
  48. """Convert the given result data to a string representation.
  49. Args:
  50. json_format (bool): If True, return a JSON formatted string. Default is False.
  51. indent (int): Number of spaces to indent for JSON formatting. Default is 4.
  52. ensure_ascii (bool): If True, ensure all characters are ASCII. Default is False.
  53. Returns:
  54. Dict[str, str]: The string representation of the result.
  55. """
  56. return {"res": str(self)}
  57. def print(self) -> None:
  58. """Print the string representation of the result."""
  59. logging.info(self.str)
  60. class JsonMixin:
  61. """Mixin class for adding JSON serialization capabilities."""
  62. def __init__(self) -> None:
  63. self._json_writer = JsonWriter()
  64. self._save_funcs.append(self.save_to_json)
  65. def _to_json(self) -> Dict[str, Dict[str, Any]]:
  66. """Convert the object to a JSON-serializable format.
  67. Returns:
  68. Dict[str, Dict[str, Any]]: A dictionary representation of the object that is JSON-serializable.
  69. """
  70. def _format_data(obj):
  71. """Helper function to format data into a JSON-serializable format.
  72. Args:
  73. obj: The object to be formatted.
  74. Returns:
  75. Any: The formatted object.
  76. """
  77. if isinstance(obj, np.float32):
  78. return float(obj)
  79. elif isinstance(obj, np.ndarray):
  80. return [_format_data(item) for item in obj.tolist()]
  81. elif isinstance(obj, pd.DataFrame):
  82. return json.loads(obj.to_json(orient="records", force_ascii=False))
  83. elif isinstance(obj, Path):
  84. return obj.as_posix()
  85. elif isinstance(obj, dict):
  86. return dict({k: _format_data(v) for k, v in obj.items()})
  87. elif isinstance(obj, (list, tuple)):
  88. return [_format_data(i) for i in obj]
  89. else:
  90. return obj
  91. return {"res": _format_data(copy.deepcopy(self))}
  92. @property
  93. def json(self) -> Dict[str, Dict[str, Any]]:
  94. """Property to get the JSON representation of the result.
  95. Returns:
  96. Dict[str, Dict[str, Any]]: The dict type JSON representation of the result.
  97. """
  98. return self._to_json()
  99. def save_to_json(
  100. self,
  101. save_path: str,
  102. indent: int = 4,
  103. ensure_ascii: bool = False,
  104. *args: List,
  105. **kwargs: Dict,
  106. ) -> None:
  107. """Save the JSON representation of the object to a file.
  108. Args:
  109. save_path (str): The path to save the JSON file. If the save path does not end with '.json', it appends the base name and suffix of the input path.
  110. indent (int): The number of spaces to indent for pretty printing. Default is 4.
  111. ensure_ascii (bool): If False, non-ASCII characters will be included in the output. Default is False.
  112. *args: Additional positional arguments to pass to the underlying writer.
  113. **kwargs: Additional keyword arguments to pass to the underlying writer.
  114. """
  115. def _is_json_file(file_path):
  116. mime_type, _ = mimetypes.guess_type(file_path)
  117. return mime_type is not None and mime_type == "application/json"
  118. if not _is_json_file(save_path):
  119. fp = Path(self["input_path"])
  120. stem = fp.stem
  121. suffix = fp.suffix
  122. base_save_path = Path(save_path)
  123. for key in self.json:
  124. save_path = base_save_path / f"{stem}_{key}.json"
  125. self._json_writer.write(
  126. save_path.as_posix(), self.json[key], *args, **kwargs
  127. )
  128. else:
  129. if len(self.json) > 1:
  130. logging.warning(
  131. f"The result has multiple json files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  132. )
  133. self._json_writer.write(
  134. save_path,
  135. self.json[list(self.json.keys())[0]],
  136. indent=indent,
  137. ensure_ascii=ensure_ascii,
  138. *args,
  139. **kwargs,
  140. )
  141. def print(
  142. self, json_format: bool = False, indent: int = 4, ensure_ascii: bool = False
  143. ) -> None:
  144. """Print the string representation of the result.
  145. Args:
  146. json_format (bool): If True, print a JSON formatted string. Default is False.
  147. indent (int): Number of spaces to indent for JSON formatting. Default is 4.
  148. ensure_ascii (bool): If True, ensure all characters are ASCII. Default is False.
  149. """
  150. if json_format:
  151. str_ = json.dumps(self.json, indent=indent, ensure_ascii=ensure_ascii)
  152. else:
  153. str_ = str(self.json)
  154. logging.info(str_)
  155. class Base64Mixin:
  156. """Mixin class for adding Base64 encoding capabilities."""
  157. def __init__(self, *args: List, **kwargs: Dict) -> None:
  158. """Initializes the Base64Mixin.
  159. Args:
  160. *args: Positional arguments to pass to the TextWriter.
  161. **kwargs: Keyword arguments to pass to the TextWriter.
  162. """
  163. self._base64_writer = TextWriter(*args, **kwargs)
  164. self._save_funcs.append(self.save_to_base64)
  165. @abstractmethod
  166. def _to_base64(self) -> Dict[str, str]:
  167. """Abstract method to convert the result to Base64.
  168. Returns:
  169. Dict[str, str]: The str type Base64 representation result.
  170. """
  171. raise NotImplementedError
  172. @property
  173. def base64(self) -> Dict[str, str]:
  174. """
  175. Property that returns the Base64 encoded content.
  176. Returns:
  177. Dict[str, str]: The base64 representation of the result.
  178. """
  179. return self._to_base64()
  180. def save_to_base64(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  181. """Saves the Base64 encoded content to the specified path.
  182. Args:
  183. save_path (str): The path to save the base64 representation result. If the save path does not end with '.b64', it appends the base name and suffix of the input path.
  184. *args: Additional positional arguments that will be passed to the base64 writer.
  185. **kwargs: Additional keyword arguments that will be passed to the base64 writer.
  186. """
  187. if not str(save_path).lower().endswith((".b64")):
  188. fp = Path(self["input_path"])
  189. stem = fp.stem
  190. suffix = fp.suffix
  191. base_save_path = Path(save_path)
  192. for key in self.base64:
  193. save_path = base_save_path / f"{stem}_{key}.b64"
  194. self._base64_writer.write(
  195. save_path.as_posix(), self.base64[key], *args, **kwargs
  196. )
  197. else:
  198. if len(self.base64) > 1:
  199. logging.warning(
  200. f"The result has multiple base64 files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  201. )
  202. self._base64_writer.write(
  203. save_path, self.base64[list(self.base64.keys())[0]], *args, **kwargs
  204. )
  205. class ImgMixin:
  206. """Mixin class for adding image handling capabilities."""
  207. def __init__(self, backend: str = "pillow", *args: List, **kwargs: Dict) -> None:
  208. """Initializes ImgMixin.
  209. Args:
  210. backend (str): The backend to use for image processing. Defaults to "pillow".
  211. *args: Additional positional arguments to pass to the ImageWriter.
  212. **kwargs: Additional keyword arguments to pass to the ImageWriter.
  213. """
  214. self._img_writer = ImageWriter(backend=backend, *args, **kwargs)
  215. self._save_funcs.append(self.save_to_img)
  216. @abstractmethod
  217. def _to_img(self) -> Dict[str, Image.Image]:
  218. """Abstract method to convert the result to an image.
  219. Returns:
  220. Dict[str, Image.Image]: The image representation result.
  221. """
  222. raise NotImplementedError
  223. @property
  224. def img(self) -> Dict[str, Image.Image]:
  225. """Property to get the image representation of the result.
  226. Returns:
  227. Dict[str, Image.Image]: The image representation of the result.
  228. """
  229. return self._to_img()
  230. def save_to_img(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  231. """Saves the image representation of the result to the specified path.
  232. Args:
  233. save_path (str): The path to save the image. If the save path does not end with .jpg or .png, it appends the input path's stem and suffix to the save path.
  234. *args: Additional positional arguments that will be passed to the image writer.
  235. **kwargs: Additional keyword arguments that will be passed to the image writer.
  236. """
  237. def _is_image_file(file_path):
  238. mime_type, _ = mimetypes.guess_type(file_path)
  239. return mime_type is not None and mime_type.startswith("image/")
  240. if not _is_image_file(save_path):
  241. fp = Path(self["input_path"])
  242. stem = fp.stem
  243. suffix = fp.suffix
  244. base_save_path = Path(save_path)
  245. for key in self.img:
  246. save_path = base_save_path / f"{stem}_{key}{suffix}"
  247. self._img_writer.write(
  248. save_path.as_posix(), self.img[key], *args, **kwargs
  249. )
  250. else:
  251. if len(self.img) > 1:
  252. logging.warning(
  253. f"The result has multiple img files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  254. )
  255. self._img_writer.write(
  256. save_path, self.img[list(self.img.keys())[0]], *args, **kwargs
  257. )
  258. class CSVMixin:
  259. """Mixin class for adding CSV handling capabilities."""
  260. def __init__(self, backend: str = "pandas", *args: List, **kwargs: Dict) -> None:
  261. """Initializes the CSVMixin.
  262. Args:
  263. backend (str): The backend to use for CSV operations (default is "pandas").
  264. *args: Optional positional arguments to pass to the CSVWriter.
  265. **kwargs: Optional keyword arguments to pass to the CSVWriter.
  266. """
  267. self._csv_writer = CSVWriter(backend=backend, *args, **kwargs)
  268. if not hasattr(self, "_save_funcs"):
  269. self._save_funcs = []
  270. self._save_funcs.append(self.save_to_csv)
  271. @property
  272. def csv(self) -> Dict[str, pd.DataFrame]:
  273. """Property to get the pandas Dataframe representation of the result.
  274. Returns:
  275. Dict[str, pd.DataFrame]: The pandas.DataFrame representation of the result.
  276. """
  277. return self._to_csv()
  278. @abstractmethod
  279. def _to_csv(self) -> Dict[str, pd.DataFrame]:
  280. """Abstract method to convert the result to pandas.DataFrame.
  281. Returns:
  282. Dict[str, pd.DataFrame]: The pandas.DataFrame representation result.
  283. """
  284. raise NotImplementedError
  285. def save_to_csv(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  286. """Saves the result to a CSV file.
  287. Args:
  288. save_path (str): The path to save the CSV file. If the path does not end with ".csv",
  289. the stem of the input path attribute (self['input_path']) will be used as the filename.
  290. *args: Optional positional arguments to pass to the CSV writer's write method.
  291. **kwargs: Optional keyword arguments to pass to the CSV writer's write method.
  292. """
  293. def _is_csv_file(file_path):
  294. mime_type, _ = mimetypes.guess_type(file_path)
  295. return mime_type is not None and mime_type == "text/csv"
  296. if not _is_csv_file(save_path):
  297. fp = Path(self["input_path"])
  298. stem = fp.stem
  299. base_save_path = Path(save_path)
  300. for key in self.csv:
  301. save_path = base_save_path / f"{stem}_{key}.csv"
  302. self._csv_writer.write(
  303. save_path.as_posix(), self.csv[key], *args, **kwargs
  304. )
  305. else:
  306. if len(self.csv) > 1:
  307. logging.warning(
  308. f"The result has multiple csv files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  309. )
  310. self._csv_writer.write(
  311. save_path, self.csv[list(self.csv.keys())[0]], *args, **kwargs
  312. )
  313. class HtmlMixin:
  314. """Mixin class for adding HTML handling capabilities."""
  315. def __init__(self, *args: List, **kwargs: Dict) -> None:
  316. """
  317. Initializes the HTML writer and appends the save_to_html method to the save functions list.
  318. Args:
  319. *args: Positional arguments passed to the HtmlWriter.
  320. **kwargs: Keyword arguments passed to the HtmlWriter.
  321. """
  322. self._html_writer = HtmlWriter(*args, **kwargs)
  323. self._save_funcs.append(self.save_to_html)
  324. @property
  325. def html(self) -> Dict[str, str]:
  326. """Property to get the HTML representation of the result.
  327. Returns:
  328. str: The str type HTML representation of the result.
  329. """
  330. return self._to_html()
  331. @abstractmethod
  332. def _to_html(self) -> Dict[str, str]:
  333. """Abstract method to convert the result to str type HTML representation.
  334. Returns:
  335. Dict[str, str]: The str type HTML representation result.
  336. """
  337. raise NotImplementedError
  338. def save_to_html(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  339. """Saves the HTML representation of the object to the specified path.
  340. Args:
  341. save_path (str): The path to save the HTML file.
  342. *args: Additional positional arguments.
  343. **kwargs: Additional keyword arguments.
  344. """
  345. def _is_html_file(file_path):
  346. mime_type, _ = mimetypes.guess_type(file_path)
  347. return mime_type is not None and mime_type == "text/html"
  348. if not _is_html_file(save_path):
  349. fp = Path(self["input_path"])
  350. stem = fp.stem
  351. base_save_path = Path(save_path)
  352. for key in self.html:
  353. save_path = base_save_path / f"{stem}_{key}.html"
  354. self._html_writer.write(
  355. save_path.as_posix(), self.html[key], *args, **kwargs
  356. )
  357. else:
  358. if len(self.html) > 1:
  359. logging.warning(
  360. f"The result has multiple html files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  361. )
  362. self._html_writer.write(
  363. save_path, self.html[list(self.html.keys())[0]], *args, **kwargs
  364. )
  365. class XlsxMixin:
  366. """Mixin class for adding XLSX handling capabilities."""
  367. def __init__(self, *args: List, **kwargs: Dict) -> None:
  368. """Initializes the XLSX writer and appends the save_to_xlsx method to the save functions.
  369. Args:
  370. *args: Positional arguments to be passed to the XlsxWriter constructor.
  371. **kwargs: Keyword arguments to be passed to the XlsxWriter constructor.
  372. """
  373. self._xlsx_writer = XlsxWriter(*args, **kwargs)
  374. self._save_funcs.append(self.save_to_xlsx)
  375. @property
  376. def xlsx(self) -> Dict[str, str]:
  377. """Property to get the XLSX representation of the result.
  378. Returns:
  379. Dict[str, str]: The str type XLSX representation of the result.
  380. """
  381. return self._to_xlsx()
  382. @abstractmethod
  383. def _to_xlsx(self) -> Dict[str, str]:
  384. """Abstract method to convert the result to str type XLSX representation.
  385. Returns:
  386. Dict[str, str]: The str type HTML representation result.
  387. """
  388. raise NotImplementedError
  389. def save_to_xlsx(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  390. """Saves the HTML representation to an XLSX file.
  391. Args:
  392. save_path (str): The path to save the XLSX file. If the path does not end with ".xlsx",
  393. the filename will be set to the stem of the input path with ".xlsx" extension.
  394. *args: Additional positional arguments to pass to the XLSX writer.
  395. **kwargs: Additional keyword arguments to pass to the XLSX writer.
  396. """
  397. def _is_xlsx_file(file_path):
  398. mime_type, _ = mimetypes.guess_type(file_path)
  399. return (
  400. mime_type is not None
  401. and mime_type
  402. == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
  403. )
  404. if not _is_xlsx_file(save_path):
  405. fp = Path(self["input_path"])
  406. stem = fp.stem
  407. base_save_path = Path(save_path)
  408. for key in self.xlsx:
  409. save_path = base_save_path / f"{stem}_{key}.xlsx"
  410. self._xlsx_writer.write(
  411. save_path.as_posix(), self.xlsx[key], *args, **kwargs
  412. )
  413. else:
  414. if len(self.xlsx) > 1:
  415. logging.warning(
  416. f"The result has multiple xlsx files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  417. )
  418. self._xlsx_writer.write(
  419. save_path, self.xlsx[list(self.xlsx.keys())[0]], *args, **kwargs
  420. )
  421. class VideoMixin:
  422. """Mixin class for adding Video handling capabilities."""
  423. def __init__(self, backend: str = "opencv", *args: List, **kwargs: Dict) -> None:
  424. """Initializes VideoMixin.
  425. Args:
  426. backend (str): The backend to use for video processing. Defaults to "opencv".
  427. *args: Additional positional arguments to pass to the VideoWriter.
  428. **kwargs: Additional keyword arguments to pass to the VideoWriter.
  429. """
  430. self._backend = backend
  431. self._save_funcs.append(self.save_to_video)
  432. @abstractmethod
  433. def _to_video(self) -> Dict[str, np.array]:
  434. """Abstract method to convert the result to a video.
  435. Returns:
  436. Dict[str, np.array]: The video representation result.
  437. """
  438. raise NotImplementedError
  439. @property
  440. def video(self) -> Dict[str, np.array]:
  441. """Property to get the video representation of the result.
  442. Returns:
  443. Dict[str, np.array]: The video representation of the result.
  444. """
  445. return self._to_video()
  446. def save_to_video(self, save_path: str, *args: List, **kwargs: Dict) -> None:
  447. """Saves the video representation of the result to the specified path.
  448. Args:
  449. save_path (str): The path to save the video. If the save path does not end with .mp4 or .avi, it appends the input path's stem and suffix to the save path.
  450. *args: Additional positional arguments that will be passed to the video writer.
  451. **kwargs: Additional keyword arguments that will be passed to the video writer.
  452. """
  453. def _is_video_file(file_path):
  454. mime_type, _ = mimetypes.guess_type(file_path)
  455. return mime_type is not None and mime_type.startswith("video/")
  456. video_writer = VideoWriter(backend=self._backend, *args, **kwargs)
  457. if not _is_video_file(save_path):
  458. fp = Path(self["input_path"])
  459. stem = fp.stem
  460. suffix = fp.suffix
  461. base_save_path = Path(save_path)
  462. for key in self.video:
  463. save_path = base_save_path / f"{stem}_{key}{suffix}"
  464. video_writer.write(
  465. save_path.as_posix(), self.video[key], *args, **kwargs
  466. )
  467. else:
  468. if len(self.video) > 1:
  469. logging.warning(
  470. f"The result has multiple video files need to be saved. But the `save_path` has been specfied as `{save_path}`!"
  471. )
  472. video_writer.write(
  473. save_path, self.video[list(self.video.keys())[0]], *args, **kwargs
  474. )
  475. class MarkdownMixin:
  476. def __init__(self, *args: list, **kwargs: dict):
  477. self._markdown_writer = MarkdownWriter(*args, **kwargs)
  478. self._save_funcs.append(self.save_to_markdown)
  479. @abstractmethod
  480. def _to_markdown(self):
  481. """
  482. Convert the result to markdown format.
  483. Returns:
  484. Dict
  485. """
  486. raise NotImplementedError
  487. @property
  488. def markdown(self):
  489. return self._to_markdown()
  490. def save_to_markdown(self, save_path, *args, **kwargs):
  491. save_path = Path(save_path)
  492. if not save_path.suffix.lower() == ".md":
  493. save_path = save_path / f"layout_parsing_result.md"
  494. self.save_path = save_path
  495. self._save_list_data(
  496. self._markdown_writer.write,
  497. save_path,
  498. self.markdown,
  499. *args,
  500. **kwargs,
  501. )
  502. def _save_list_data(self, save_func, save_path, data, *args, **kwargs):
  503. save_path = Path(save_path)
  504. if data is None:
  505. return
  506. if isinstance(data, list):
  507. for idx, single in enumerate(data):
  508. save_func(
  509. (
  510. save_path.parent / f"{save_path.stem}_{idx}{save_path.suffix}"
  511. ).as_posix(),
  512. single,
  513. *args,
  514. **kwargs,
  515. )
  516. save_func(save_path.as_posix(), data, *args, **kwargs)
  517. logging.info(f"The result has been saved in {save_path}.")