writers.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. import enum
  15. import json
  16. from pathlib import Path
  17. import cv2
  18. import numpy as np
  19. import pandas as pd
  20. import yaml
  21. from PIL import Image
  22. from .tablepyxl import document_to_xl
  23. __all__ = [
  24. "WriterType",
  25. "ImageWriter",
  26. "TextWriter",
  27. "JsonWriter",
  28. "CSVWriter",
  29. "HtmlWriter",
  30. "XlsxWriter",
  31. "YAMLWriter",
  32. "VideoWriter",
  33. "MarkdownWriter",
  34. ]
  35. class WriterType(enum.Enum):
  36. """WriterType"""
  37. IMAGE = 1
  38. VIDEO = 2
  39. TEXT = 3
  40. JSON = 4
  41. HTML = 5
  42. XLSX = 6
  43. CSV = 7
  44. YAML = 8
  45. class _BaseWriter(object):
  46. """_BaseWriter"""
  47. def __init__(self, backend, **bk_args):
  48. super().__init__()
  49. if len(bk_args) == 0:
  50. bk_args = self.get_default_backend_args()
  51. self.bk_type = backend
  52. self.bk_args = bk_args
  53. self._backend = self.get_backend()
  54. def write(self, out_path, obj):
  55. """write"""
  56. raise NotImplementedError
  57. def get_backend(self, bk_args=None):
  58. """get backend"""
  59. if bk_args is None:
  60. bk_args = self.bk_args
  61. return self._init_backend(self.bk_type, bk_args)
  62. def set_backend(self, backend, **bk_args):
  63. self.bk_type = backend
  64. self.bk_args = bk_args
  65. self._backend = self.get_backend()
  66. def _init_backend(self, bk_type, bk_args):
  67. """init backend"""
  68. raise NotImplementedError
  69. def get_type(self):
  70. """get type"""
  71. raise NotImplementedError
  72. def get_default_backend_args(self):
  73. """get default backend arguments"""
  74. return {}
  75. class ImageWriter(_BaseWriter):
  76. """ImageWriter"""
  77. def __init__(self, backend="opencv", **bk_args):
  78. super().__init__(backend=backend, **bk_args)
  79. def write(self, out_path, obj):
  80. """write"""
  81. return self._backend.write_obj(str(out_path), obj)
  82. def _init_backend(self, bk_type, bk_args):
  83. """init backend"""
  84. if bk_type == "opencv":
  85. return OpenCVImageWriterBackend(**bk_args)
  86. elif bk_type == "pil" or bk_type == "pillow":
  87. return PILImageWriterBackend(**bk_args)
  88. else:
  89. raise ValueError("Unsupported backend type")
  90. def get_type(self):
  91. """get type"""
  92. return WriterType.IMAGE
  93. class VideoWriter(_BaseWriter):
  94. """VideoWriter"""
  95. def __init__(self, backend="opencv", **bk_args):
  96. super().__init__(backend=backend, **bk_args)
  97. def write(self, out_path, obj):
  98. """write"""
  99. return self._backend.write_obj(str(out_path), obj)
  100. def _init_backend(self, bk_type, bk_args):
  101. """init backend"""
  102. if bk_type == "opencv":
  103. return OpenCVVideoWriterBackend(**bk_args)
  104. else:
  105. raise ValueError("Unsupported backend type")
  106. def get_type(self):
  107. """get type"""
  108. return WriterType.VIDEO
  109. class TextWriter(_BaseWriter):
  110. """TextWriter"""
  111. def __init__(self, backend="python", **bk_args):
  112. super().__init__(backend=backend, **bk_args)
  113. def write(self, out_path, obj):
  114. """write"""
  115. return self._backend.write_obj(str(out_path), obj)
  116. def _init_backend(self, bk_type, bk_args):
  117. """init backend"""
  118. if bk_type == "python":
  119. return TextWriterBackend(**bk_args)
  120. else:
  121. raise ValueError("Unsupported backend type")
  122. def get_type(self):
  123. """get type"""
  124. return WriterType.TEXT
  125. class JsonWriter(_BaseWriter):
  126. def __init__(self, backend="json", **bk_args):
  127. super().__init__(backend=backend, **bk_args)
  128. def write(self, out_path, obj, **bk_args):
  129. return self._backend.write_obj(str(out_path), obj, **bk_args)
  130. def _init_backend(self, bk_type, bk_args):
  131. if bk_type == "json":
  132. return JsonWriterBackend(**bk_args)
  133. elif bk_type == "ujson":
  134. return UJsonWriterBackend(**bk_args)
  135. else:
  136. raise ValueError("Unsupported backend type")
  137. def get_type(self):
  138. """get type"""
  139. return WriterType.JSON
  140. class HtmlWriter(_BaseWriter):
  141. def __init__(self, backend="html", **bk_args):
  142. super().__init__(backend=backend, **bk_args)
  143. def write(self, out_path, obj, **bk_args):
  144. return self._backend.write_obj(str(out_path), obj, **bk_args)
  145. def _init_backend(self, bk_type, bk_args):
  146. if bk_type == "html":
  147. return HtmlWriterBackend(**bk_args)
  148. else:
  149. raise ValueError("Unsupported backend type")
  150. def get_type(self):
  151. """get type"""
  152. return WriterType.HTML
  153. class XlsxWriter(_BaseWriter):
  154. def __init__(self, backend="xlsx", **bk_args):
  155. super().__init__(backend=backend, **bk_args)
  156. def write(self, out_path, obj, **bk_args):
  157. return self._backend.write_obj(str(out_path), obj, **bk_args)
  158. def _init_backend(self, bk_type, bk_args):
  159. if bk_type == "xlsx":
  160. return XlsxWriterBackend(**bk_args)
  161. else:
  162. raise ValueError("Unsupported backend type")
  163. def get_type(self):
  164. """get type"""
  165. return WriterType.XLSX
  166. class YAMLWriter(_BaseWriter):
  167. def __init__(self, backend="PyYAML", **bk_args):
  168. super().__init__(backend=backend, **bk_args)
  169. def write(self, out_path, obj, **bk_args):
  170. return self._backend.write_obj(str(out_path), obj, **bk_args)
  171. def _init_backend(self, bk_type, bk_args):
  172. if bk_type == "PyYAML":
  173. return YAMLWriterBackend(**bk_args)
  174. else:
  175. raise ValueError("Unsupported backend type")
  176. def get_type(self):
  177. """get type"""
  178. return WriterType.YAML
  179. class MarkdownWriter(_BaseWriter):
  180. """MarkdownWriter"""
  181. def __init__(self, backend="markdown", **bk_args):
  182. super().__init__(backend=backend, **bk_args)
  183. def write(self, out_path, obj):
  184. """write"""
  185. return self._backend.write_obj(str(out_path), obj)
  186. def _init_backend(self, bk_type, bk_args):
  187. """init backend"""
  188. if bk_type == "markdown":
  189. return MarkdownWriterBackend(**bk_args)
  190. else:
  191. raise ValueError("Unsupported backend type")
  192. def get_type(self):
  193. """get type"""
  194. return WriterType.MARKDOWN
  195. class _BaseWriterBackend(object):
  196. """_BaseWriterBackend"""
  197. def write_obj(self, out_path, obj, **bk_args):
  198. """write object"""
  199. Path(out_path).parent.mkdir(parents=True, exist_ok=True)
  200. return self._write_obj(out_path, obj, **bk_args)
  201. def _write_obj(self, out_path, obj, **bk_args):
  202. """write object"""
  203. raise NotImplementedError
  204. class TextWriterBackend(_BaseWriterBackend):
  205. """TextWriterBackend"""
  206. def __init__(self, mode="w", encoding="utf-8"):
  207. super().__init__()
  208. self.mode = mode
  209. self.encoding = encoding
  210. def _write_obj(self, out_path, obj):
  211. """write text object"""
  212. with open(out_path, mode=self.mode, encoding=self.encoding) as f:
  213. f.write(obj)
  214. class HtmlWriterBackend(_BaseWriterBackend):
  215. def __init__(self, mode="w", encoding="utf-8"):
  216. super().__init__()
  217. self.mode = mode
  218. self.encoding = encoding
  219. def _write_obj(self, out_path, obj, **bk_args):
  220. with open(out_path, mode=self.mode, encoding=self.encoding) as f:
  221. f.write(obj)
  222. class XlsxWriterBackend(_BaseWriterBackend):
  223. def _write_obj(self, out_path, obj, **bk_args):
  224. document_to_xl(obj, out_path)
  225. class _ImageWriterBackend(_BaseWriterBackend):
  226. """_ImageWriterBackend"""
  227. class OpenCVImageWriterBackend(_ImageWriterBackend):
  228. """OpenCVImageWriterBackend"""
  229. def _write_obj(self, out_path, obj):
  230. """write image object by OpenCV"""
  231. if isinstance(obj, Image.Image):
  232. # Assuming the channel order is RGB.
  233. arr = np.asarray(obj)[:, :, ::-1]
  234. elif isinstance(obj, np.ndarray):
  235. arr = obj
  236. else:
  237. raise TypeError("Unsupported object type")
  238. return cv2.imwrite(out_path, arr)
  239. class PILImageWriterBackend(_ImageWriterBackend):
  240. """PILImageWriterBackend"""
  241. def __init__(self, format_=None):
  242. super().__init__()
  243. self.format = format_
  244. def _write_obj(self, out_path, obj):
  245. """write image object by PIL"""
  246. if isinstance(obj, Image.Image):
  247. img = obj
  248. elif isinstance(obj, np.ndarray):
  249. img = Image.fromarray(obj)
  250. else:
  251. raise TypeError("Unsupported object type")
  252. if len(img.getbands()) == 4:
  253. self.format = "PNG"
  254. return img.save(out_path, format=self.format)
  255. class _VideoWriterBackend(_BaseWriterBackend):
  256. """_VideoWriterBackend"""
  257. class OpenCVVideoWriterBackend(_VideoWriterBackend):
  258. """OpenCVImageWriterBackend"""
  259. def _write_obj(self, out_path, obj):
  260. """write video object by OpenCV"""
  261. obj, fps = obj
  262. if isinstance(obj, np.ndarray):
  263. vr = obj
  264. width, height = vr[0].shape[1], vr[0].shape[0]
  265. fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Alternatively, use 'XVID'
  266. out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
  267. for frame in vr:
  268. out.write(frame)
  269. out.release()
  270. else:
  271. raise TypeError("Unsupported object type")
  272. class _BaseJsonWriterBackend(object):
  273. def __init__(self, indent=4, ensure_ascii=False):
  274. super().__init__()
  275. self.indent = indent
  276. self.ensure_ascii = ensure_ascii
  277. def write_obj(self, out_path, obj, **bk_args):
  278. Path(out_path).parent.mkdir(parents=True, exist_ok=True)
  279. return self._write_obj(out_path, obj, **bk_args)
  280. def _write_obj(self, out_path, obj):
  281. raise NotImplementedError
  282. class JsonWriterBackend(_BaseJsonWriterBackend):
  283. def _write_obj(self, out_path, obj, **bk_args):
  284. with open(out_path, "w", encoding="utf-8") as f:
  285. json.dump(obj, f, **bk_args)
  286. class UJsonWriterBackend(_BaseJsonWriterBackend):
  287. # TODO
  288. def _write_obj(self, out_path, obj, **bk_args):
  289. raise NotImplementedError
  290. class YAMLWriterBackend(_BaseWriterBackend):
  291. def __init__(self, mode="w", encoding="utf-8"):
  292. super().__init__()
  293. self.mode = mode
  294. self.encoding = encoding
  295. def _write_obj(self, out_path, obj, **bk_args):
  296. """write text object"""
  297. with open(out_path, mode=self.mode, encoding=self.encoding) as f:
  298. yaml.dump(obj, f, **bk_args)
  299. class CSVWriter(_BaseWriter):
  300. """CSVWriter"""
  301. def __init__(self, backend="pandas", **bk_args):
  302. super().__init__(backend=backend, **bk_args)
  303. def write(self, out_path, obj):
  304. """write"""
  305. return self._backend.write_obj(str(out_path), obj)
  306. def _init_backend(self, bk_type, bk_args):
  307. """init backend"""
  308. if bk_type == "pandas":
  309. return PandasCSVWriterBackend(**bk_args)
  310. else:
  311. raise ValueError("Unsupported backend type")
  312. def get_type(self):
  313. """get type"""
  314. return WriterType.CSV
  315. class _CSVWriterBackend(_BaseWriterBackend):
  316. """_CSVWriterBackend"""
  317. class PandasCSVWriterBackend(_CSVWriterBackend):
  318. """PILImageWriterBackend"""
  319. def __init__(self):
  320. super().__init__()
  321. def _write_obj(self, out_path, obj):
  322. """write image object by PIL"""
  323. if isinstance(obj, pd.DataFrame):
  324. ts = obj
  325. else:
  326. raise TypeError("Unsupported object type")
  327. return ts.to_csv(out_path)
  328. class MarkdownWriterBackend(_BaseWriterBackend):
  329. """MarkdownWriterBackend"""
  330. def __init__(self):
  331. super().__init__()
  332. def _write_obj(self, out_path, obj):
  333. """write markdown obj"""
  334. with open(out_path, mode="w", encoding="utf-8", errors="replace") as f:
  335. f.write(obj)