writers.py 11 KB

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