readers.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433
  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 enum
  15. import itertools
  16. import cv2
  17. import fitz
  18. from PIL import Image, ImageOps
  19. import pandas as pd
  20. import numpy as np
  21. import yaml
  22. import decord
  23. import random
  24. __all__ = [
  25. "ReaderType",
  26. "ImageReader",
  27. "VideoReader",
  28. "CSVReader",
  29. "PDFReader",
  30. "YAMLReader",
  31. ]
  32. class ReaderType(enum.Enum):
  33. """ReaderType"""
  34. IMAGE = 1
  35. GENERATIVE = 2
  36. POINT_CLOUD = 3
  37. JSON = 4
  38. TS = 5
  39. PDF = 6
  40. YAML = 8
  41. class _BaseReader(object):
  42. """_BaseReader"""
  43. def __init__(self, backend, **bk_args):
  44. super().__init__()
  45. if len(bk_args) == 0:
  46. bk_args = self.get_default_backend_args()
  47. self.bk_type = backend
  48. self.bk_args = bk_args
  49. self._backend = self.get_backend()
  50. def read(self, in_path):
  51. """read file from path"""
  52. raise NotImplementedError
  53. def get_backend(self, bk_args=None):
  54. """get the backend"""
  55. if bk_args is None:
  56. bk_args = self.bk_args
  57. return self._init_backend(self.bk_type, bk_args)
  58. def set_backend(self, backend, **bk_args):
  59. self.bk_type = backend
  60. self.bk_args = bk_args
  61. self._backend = self.get_backend()
  62. def _init_backend(self, bk_type, bk_args):
  63. """init backend"""
  64. raise NotImplementedError
  65. def get_type(self):
  66. """get type"""
  67. raise NotImplementedError
  68. def get_default_backend_args(self):
  69. """get default backend arguments"""
  70. return {}
  71. class PDFReader(_BaseReader):
  72. """PDFReader"""
  73. def __init__(self, backend="fitz", **bk_args):
  74. super().__init__(backend, **bk_args)
  75. def read(self, in_path):
  76. yield from self._backend.read_file(str(in_path))
  77. def _init_backend(self, bk_type, bk_args):
  78. return PDFReaderBackend(**bk_args)
  79. def get_type(self):
  80. return ReaderType.PDF
  81. class ImageReader(_BaseReader):
  82. """ImageReader"""
  83. def __init__(self, backend="opencv", **bk_args):
  84. super().__init__(backend=backend, **bk_args)
  85. def read(self, in_path):
  86. """read the image file from path"""
  87. arr = self._backend.read_file(str(in_path))
  88. return arr
  89. def _init_backend(self, bk_type, bk_args):
  90. """init backend"""
  91. if bk_type == "opencv":
  92. return OpenCVImageReaderBackend(**bk_args)
  93. elif bk_type == "pil" or bk_type == "pillow":
  94. return PILImageReaderBackend(**bk_args)
  95. else:
  96. raise ValueError("Unsupported backend type")
  97. def get_type(self):
  98. """get type"""
  99. return ReaderType.IMAGE
  100. class _GenerativeReader(_BaseReader):
  101. """_GenerativeReader"""
  102. def get_type(self):
  103. """get type"""
  104. return ReaderType.GENERATIVE
  105. def is_generative_reader(reader):
  106. """is_generative_reader"""
  107. return isinstance(reader, _GenerativeReader)
  108. class VideoReader(_GenerativeReader):
  109. """VideoReader"""
  110. def __init__(
  111. self,
  112. backend="opencv",
  113. st_frame_id=0,
  114. max_num_frames=None,
  115. auto_close=True,
  116. **bk_args,
  117. ):
  118. super().__init__(backend=backend, **bk_args)
  119. self.st_frame_id = st_frame_id
  120. self.max_num_frames = max_num_frames
  121. self.auto_close = auto_close
  122. self._fps = 0
  123. def read(self, in_path):
  124. """read vide file from path"""
  125. self._backend.set_pos(self.st_frame_id)
  126. gen = self._backend.read_file(str(in_path))
  127. if self.max_num_frames is not None:
  128. gen = itertools.islice(gen, self.num_frames)
  129. yield from gen
  130. if self.auto_close:
  131. self._backend.close()
  132. def get_fps(self):
  133. """get fps"""
  134. return self._backend.get_fps()
  135. def _init_backend(self, bk_type, bk_args):
  136. """init backend"""
  137. if bk_type == "opencv":
  138. return OpenCVVideoReaderBackend(**bk_args)
  139. elif bk_type == "decord":
  140. return DecordVideoReaderBackend(**bk_args)
  141. else:
  142. raise ValueError("Unsupported backend type")
  143. class YAMLReader(_BaseReader):
  144. def __init__(self, backend="PyYAML", **bk_args):
  145. super().__init__(backend, **bk_args)
  146. def read(self, in_path):
  147. return self._backend.read_file(str(in_path))
  148. def _init_backend(self, bk_type, bk_args):
  149. if bk_type == "PyYAML":
  150. return YAMLReaderBackend(**bk_args)
  151. else:
  152. raise ValueError("Unsupported backend type")
  153. def get_type(self):
  154. return ReaderType.YAML
  155. class _BaseReaderBackend(object):
  156. """_BaseReaderBackend"""
  157. def read_file(self, in_path):
  158. """read file from path"""
  159. raise NotImplementedError
  160. class _ImageReaderBackend(_BaseReaderBackend):
  161. """_ImageReaderBackend"""
  162. pass
  163. class OpenCVImageReaderBackend(_ImageReaderBackend):
  164. """OpenCVImageReaderBackend"""
  165. def __init__(self, flags=cv2.IMREAD_COLOR):
  166. super().__init__()
  167. self.flags = flags
  168. def read_file(self, in_path):
  169. """read image file from path by OpenCV"""
  170. return cv2.imread(in_path, flags=self.flags)
  171. class PILImageReaderBackend(_ImageReaderBackend):
  172. """PILImageReaderBackend"""
  173. def __init__(self):
  174. super().__init__()
  175. def read_file(self, in_path):
  176. """read image file from path by PIL"""
  177. return ImageOps.exif_transpose(Image.open(in_path))
  178. class PDFReaderBackend(_BaseReaderBackend):
  179. def __init__(self, rotate=0, zoom_x=2.0, zoom_y=2.0):
  180. super().__init__()
  181. self.mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
  182. def read_file(self, in_path):
  183. for page in fitz.open(in_path):
  184. pix = page.get_pixmap(matrix=self.mat, alpha=False)
  185. getpngdata = pix.tobytes(output="png")
  186. # decode as np.uint8
  187. image_array = np.frombuffer(getpngdata, dtype=np.uint8)
  188. img_cv = cv2.imdecode(image_array, cv2.IMREAD_ANYCOLOR)
  189. yield img_cv
  190. class _VideoReaderBackend(_BaseReaderBackend):
  191. """_VideoReaderBackend"""
  192. def set_pos(self, pos):
  193. """set pos"""
  194. raise NotImplementedError
  195. def close(self):
  196. """close io"""
  197. raise NotImplementedError
  198. class OpenCVVideoReaderBackend(_VideoReaderBackend):
  199. """OpenCVVideoReaderBackend"""
  200. def __init__(self, **bk_args):
  201. super().__init__()
  202. self.cap_init_args = bk_args
  203. self._cap = None
  204. self._pos = 0
  205. self._max_num_frames = None
  206. def get_fps(self):
  207. return self._cap.get(cv2.CAP_PROP_FPS)
  208. def read_file(self, in_path):
  209. """read vidio file from path"""
  210. if self._cap is not None:
  211. self._cap_release()
  212. self._cap = self._cap_open(in_path)
  213. if self._pos is not None:
  214. self._cap_set_pos()
  215. return self._read_frames(self._cap)
  216. def _read_frames(self, cap):
  217. """read frames"""
  218. while True:
  219. ret, frame = cap.read()
  220. if not ret:
  221. break
  222. yield frame
  223. self._cap_release()
  224. def _cap_open(self, video_path):
  225. self.cap_init_args.pop("num_seg")
  226. self.cap_init_args.pop("seg_len")
  227. self.cap_init_args.pop("sample_type")
  228. self._cap = cv2.VideoCapture(video_path, **self.cap_init_args)
  229. if not self._cap.isOpened():
  230. raise RuntimeError(f"Failed to open {video_path}")
  231. return self._cap
  232. def _cap_release(self):
  233. self._cap.release()
  234. def _cap_set_pos(self):
  235. self._cap.set(cv2.CAP_PROP_POS_FRAMES, self._pos)
  236. def set_pos(self, pos):
  237. self._pos = pos
  238. def close(self):
  239. if self._cap is not None:
  240. self._cap_release()
  241. self._cap = None
  242. class DecordVideoReaderBackend(_VideoReaderBackend):
  243. """DecordVideoReaderBackend"""
  244. def __init__(self, **bk_args):
  245. super().__init__()
  246. self.cap_init_args = bk_args
  247. self._cap = None
  248. self._pos = 0
  249. self._max_num_frames = None
  250. self.num_seg = bk_args.get("num_seg", 8)
  251. self.seg_len = bk_args.get("seg_len", 1)
  252. self.sample_type = bk_args.get("sample_type", 1)
  253. self.valid_mode = True
  254. self._fps = 0
  255. def set_pos(self, pos):
  256. self._pos = pos
  257. def sample(self, frames_len, video_object):
  258. frames_idx = []
  259. average_dur = int(frames_len / self.num_seg)
  260. for i in range(self.num_seg):
  261. idx = 0
  262. if not self.valid_mode:
  263. if average_dur >= self.seg_len:
  264. idx = random.randint(0, average_dur - self.seg_len)
  265. idx += i * average_dur
  266. elif average_dur >= 1:
  267. idx += i * average_dur
  268. else:
  269. idx = i
  270. else:
  271. if average_dur >= self.seg_len:
  272. idx = (average_dur - 1) // 2
  273. idx += i * average_dur
  274. elif average_dur >= 1:
  275. idx += i * average_dur
  276. else:
  277. idx = i
  278. for jj in range(idx, idx + self.seg_len):
  279. frames_idx.append(int(jj % frames_len))
  280. frames_select = video_object.get_batch(frames_idx)
  281. # dearray_to_img
  282. np_frames = frames_select.asnumpy()
  283. imgs = []
  284. for i in range(np_frames.shape[0]):
  285. imgbuf = np_frames[i]
  286. imgs.append(imgbuf)
  287. return imgs
  288. def get_fps(self):
  289. return self._cap.get_avg_fps()
  290. def read_file(self, in_path):
  291. """read vidio file from path"""
  292. self._cap = decord.VideoReader(in_path)
  293. frame_len = len(self._cap)
  294. if self.sample_type == "uniform":
  295. sample_video = self.sample(frame_len, self._cap)
  296. return sample_video
  297. else:
  298. return self._cap
  299. def close(self):
  300. pass
  301. class CSVReader(_BaseReader):
  302. """CSVReader"""
  303. def __init__(self, backend="pandas", **bk_args):
  304. super().__init__(backend=backend, **bk_args)
  305. def read(self, in_path):
  306. """read the image file from path"""
  307. arr = self._backend.read_file(str(in_path))
  308. return arr
  309. def _init_backend(self, bk_type, bk_args):
  310. """init backend"""
  311. if bk_type == "pandas":
  312. return PandasCSVReaderBackend(**bk_args)
  313. else:
  314. raise ValueError("Unsupported backend type")
  315. def get_type(self):
  316. """get type"""
  317. return ReaderType.TS
  318. class _CSVReaderBackend(_BaseReaderBackend):
  319. """_CSVReaderBackend"""
  320. pass
  321. class PandasCSVReaderBackend(_CSVReaderBackend):
  322. """PandasCSVReaderBackend"""
  323. def __init__(self):
  324. super().__init__()
  325. def read_file(self, in_path):
  326. """read image file from path by OpenCV"""
  327. return pd.read_csv(in_path)
  328. class YAMLReaderBackend(_BaseReaderBackend):
  329. def read_file(self, in_path, **kwargs):
  330. with open(in_path, "r", encoding="utf-8", **kwargs) as yaml_file:
  331. data = yaml.load(yaml_file, Loader=yaml.FullLoader)
  332. return data