readers.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492
  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 itertools
  16. import random
  17. import cv2
  18. import fitz
  19. import numpy as np
  20. import pandas as pd
  21. import soundfile
  22. import yaml
  23. from PIL import Image, ImageOps
  24. __all__ = [
  25. "ReaderType",
  26. "ImageReader",
  27. "VideoReader",
  28. "CSVReader",
  29. "PDFReader",
  30. "YAMLReader",
  31. "AudioReader",
  32. ]
  33. class ReaderType(enum.Enum):
  34. """ReaderType"""
  35. IMAGE = 1
  36. GENERATIVE = 2
  37. POINT_CLOUD = 3
  38. JSON = 4
  39. TS = 5
  40. PDF = 6
  41. YAML = 8
  42. class _BaseReader(object):
  43. """_BaseReader"""
  44. def __init__(self, backend, **bk_args):
  45. super().__init__()
  46. if len(bk_args) == 0:
  47. bk_args = self.get_default_backend_args()
  48. self.bk_type = backend
  49. self.bk_args = bk_args
  50. self._backend = self.get_backend()
  51. def read(self, in_path):
  52. """read file from path"""
  53. raise NotImplementedError
  54. def get_backend(self, bk_args=None):
  55. """get the backend"""
  56. if bk_args is None:
  57. bk_args = self.bk_args
  58. return self._init_backend(self.bk_type, bk_args)
  59. def set_backend(self, backend, **bk_args):
  60. self.bk_type = backend
  61. self.bk_args = bk_args
  62. self._backend = self.get_backend()
  63. def _init_backend(self, bk_type, bk_args):
  64. """init backend"""
  65. raise NotImplementedError
  66. def get_type(self):
  67. """get type"""
  68. raise NotImplementedError
  69. def get_default_backend_args(self):
  70. """get default backend arguments"""
  71. return {}
  72. class PDFReader(_BaseReader):
  73. """PDFReader"""
  74. def __init__(self, backend="fitz", **bk_args):
  75. super().__init__(backend, **bk_args)
  76. def read(self, in_path):
  77. yield from self._backend.read_file(str(in_path))
  78. def _init_backend(self, bk_type, bk_args):
  79. return PDFReaderBackend(**bk_args)
  80. def get_type(self):
  81. return ReaderType.PDF
  82. class ImageReader(_BaseReader):
  83. """ImageReader"""
  84. def __init__(self, backend="opencv", **bk_args):
  85. super().__init__(backend=backend, **bk_args)
  86. def read(self, in_path):
  87. """read the image file from path"""
  88. arr = self._backend.read_file(str(in_path))
  89. return arr
  90. def _init_backend(self, bk_type, bk_args):
  91. """init backend"""
  92. if bk_type == "opencv":
  93. return OpenCVImageReaderBackend(**bk_args)
  94. elif bk_type == "pil" or bk_type == "pillow":
  95. return PILImageReaderBackend(**bk_args)
  96. else:
  97. raise ValueError("Unsupported backend type")
  98. def get_type(self):
  99. """get type"""
  100. return ReaderType.IMAGE
  101. class _GenerativeReader(_BaseReader):
  102. """_GenerativeReader"""
  103. def get_type(self):
  104. """get type"""
  105. return ReaderType.GENERATIVE
  106. def is_generative_reader(reader):
  107. """is_generative_reader"""
  108. return isinstance(reader, _GenerativeReader)
  109. class VideoReader(_GenerativeReader):
  110. """VideoReader"""
  111. def __init__(
  112. self,
  113. backend="opencv",
  114. st_frame_id=0,
  115. max_num_frames=None,
  116. auto_close=True,
  117. **bk_args,
  118. ):
  119. super().__init__(backend=backend, **bk_args)
  120. self.st_frame_id = st_frame_id
  121. self.max_num_frames = max_num_frames
  122. self.auto_close = auto_close
  123. self._fps = 0
  124. def read(self, in_path):
  125. """read vide file from path"""
  126. self._backend.set_pos(self.st_frame_id)
  127. gen = self._backend.read_file(str(in_path))
  128. if self.max_num_frames is not None:
  129. gen = itertools.islice(gen, self.num_frames)
  130. yield from gen
  131. if self.auto_close:
  132. self._backend.close()
  133. def get_fps(self):
  134. """get fps"""
  135. return self._backend.get_fps()
  136. def _init_backend(self, bk_type, bk_args):
  137. """init backend"""
  138. if bk_type == "opencv":
  139. return OpenCVVideoReaderBackend(**bk_args)
  140. elif bk_type == "decord":
  141. return DecordVideoReaderBackend(**bk_args)
  142. else:
  143. raise ValueError("Unsupported backend type")
  144. class YAMLReader(_BaseReader):
  145. def __init__(self, backend="PyYAML", **bk_args):
  146. super().__init__(backend, **bk_args)
  147. def read(self, in_path):
  148. return self._backend.read_file(str(in_path))
  149. def _init_backend(self, bk_type, bk_args):
  150. if bk_type == "PyYAML":
  151. return YAMLReaderBackend(**bk_args)
  152. else:
  153. raise ValueError("Unsupported backend type")
  154. def get_type(self):
  155. return ReaderType.YAML
  156. class _BaseReaderBackend(object):
  157. """_BaseReaderBackend"""
  158. def read_file(self, in_path):
  159. """read file from path"""
  160. raise NotImplementedError
  161. class _ImageReaderBackend(_BaseReaderBackend):
  162. """_ImageReaderBackend"""
  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. pixmap = page.get_pixmap(matrix=self.mat, alpha=False)
  185. img_cv = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(
  186. pixmap.h, pixmap.w, pixmap.n
  187. )
  188. img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
  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.num_seg = bk_args.get("num_seg", None)
  204. self._cap = None
  205. self._pos = 0
  206. self._max_num_frames = None
  207. def get_fps(self):
  208. return self._cap.get(cv2.CAP_PROP_FPS)
  209. def read_file(self, in_path):
  210. """read video file from path"""
  211. if self._cap is not None:
  212. self._cap_release()
  213. self._cap = self._cap_open(in_path)
  214. if self._pos is not None:
  215. self._cap_set_pos()
  216. return self._read_frames(self._cap)
  217. def _read_frames(self, cap):
  218. """read frames"""
  219. if self.num_seg:
  220. queue = []
  221. while True:
  222. ret, frame = cap.read()
  223. if not ret:
  224. break
  225. queue = []
  226. if (
  227. len(queue) <= 0
  228. ): # At initialization, populate queue with initial frame
  229. for i in range(self.num_seg):
  230. queue.append(frame)
  231. queue.append(frame)
  232. queue.pop(0)
  233. yield queue.copy()
  234. else:
  235. while True:
  236. ret, frame = cap.read()
  237. if not ret:
  238. break
  239. yield frame
  240. self._cap_release()
  241. def _cap_open(self, video_path):
  242. self.cap_init_args.pop("num_seg", None)
  243. self.cap_init_args.pop("seg_len", None)
  244. self.cap_init_args.pop("sample_type", None)
  245. self._cap = cv2.VideoCapture(video_path, **self.cap_init_args)
  246. if not self._cap.isOpened():
  247. raise RuntimeError(f"Failed to open {video_path}")
  248. return self._cap
  249. def _cap_release(self):
  250. self._cap.release()
  251. def _cap_set_pos(self):
  252. self._cap.set(cv2.CAP_PROP_POS_FRAMES, self._pos)
  253. def set_pos(self, pos):
  254. self._pos = pos
  255. def close(self):
  256. if self._cap is not None:
  257. self._cap_release()
  258. self._cap = None
  259. class DecordVideoReaderBackend(_VideoReaderBackend):
  260. """DecordVideoReaderBackend"""
  261. def __init__(self, **bk_args):
  262. super().__init__()
  263. self.cap_init_args = bk_args
  264. self._cap = None
  265. self._pos = 0
  266. self._max_num_frames = None
  267. self.num_seg = bk_args.get("num_seg", 8)
  268. self.seg_len = bk_args.get("seg_len", 1)
  269. self.sample_type = bk_args.get("sample_type", 1)
  270. self.valid_mode = True
  271. self._fps = 0
  272. # XXX(gaotingquan): There is a confict with `paddle` when import `decord` globally.
  273. try:
  274. import decord
  275. self.decord_module = decord
  276. except ModuleNotFoundError():
  277. raise Exception(
  278. "Please install `decord` manually, otherwise, the related model cannot work. It can be automatically installed only on `x86_64`. Refers: `https://github.com/dmlc/decord`."
  279. )
  280. def set_pos(self, pos):
  281. self._pos = pos
  282. def sample(self, frames_len, video_object):
  283. frames_idx = []
  284. average_dur = int(frames_len / self.num_seg)
  285. for i in range(self.num_seg):
  286. idx = 0
  287. if not self.valid_mode:
  288. if average_dur >= self.seg_len:
  289. idx = random.randint(0, average_dur - self.seg_len)
  290. idx += i * average_dur
  291. elif average_dur >= 1:
  292. idx += i * average_dur
  293. else:
  294. idx = i
  295. else:
  296. if average_dur >= self.seg_len:
  297. idx = (average_dur - 1) // 2
  298. idx += i * average_dur
  299. elif average_dur >= 1:
  300. idx += i * average_dur
  301. else:
  302. idx = i
  303. for jj in range(idx, idx + self.seg_len):
  304. frames_idx.append(int(jj % frames_len))
  305. frames_select = video_object.get_batch(frames_idx)
  306. # dearray_to_img
  307. np_frames = frames_select.asnumpy()
  308. imgs = []
  309. for i in range(np_frames.shape[0]):
  310. imgbuf = np_frames[i]
  311. imgs.append(imgbuf)
  312. return imgs
  313. def get_fps(self):
  314. return self._cap.get_avg_fps()
  315. def read_file(self, in_path):
  316. """read video file from path"""
  317. self._cap = self.decord_module.VideoReader(in_path)
  318. frame_len = len(self._cap)
  319. if self.sample_type == "uniform":
  320. sample_video = self.sample(frame_len, self._cap)
  321. return sample_video
  322. else:
  323. return self._cap
  324. def close(self):
  325. pass
  326. class CSVReader(_BaseReader):
  327. """CSVReader"""
  328. def __init__(self, backend="pandas", **bk_args):
  329. super().__init__(backend=backend, **bk_args)
  330. def read(self, in_path):
  331. """read the image file from path"""
  332. arr = self._backend.read_file(str(in_path))
  333. return arr
  334. def _init_backend(self, bk_type, bk_args):
  335. """init backend"""
  336. if bk_type == "pandas":
  337. return PandasCSVReaderBackend(**bk_args)
  338. else:
  339. raise ValueError("Unsupported backend type")
  340. def get_type(self):
  341. """get type"""
  342. return ReaderType.TS
  343. class _CSVReaderBackend(_BaseReaderBackend):
  344. """_CSVReaderBackend"""
  345. class PandasCSVReaderBackend(_CSVReaderBackend):
  346. """PandasCSVReaderBackend"""
  347. def __init__(self):
  348. super().__init__()
  349. def read_file(self, in_path):
  350. """read image file from path by OpenCV"""
  351. return pd.read_csv(in_path)
  352. class YAMLReaderBackend(_BaseReaderBackend):
  353. def read_file(self, in_path, **kwargs):
  354. with open(in_path, "r", encoding="utf-8", **kwargs) as yaml_file:
  355. data = yaml.load(yaml_file, Loader=yaml.FullLoader)
  356. return data
  357. class AudioReader(_BaseReader):
  358. def __init__(self, backend="wav", **bk_args):
  359. super().__init__(backend="wav", **bk_args)
  360. def _init_backend(self, bk_type, bk_args):
  361. """init backend"""
  362. if bk_type == "wav":
  363. return WAVReaderBackend(**bk_args)
  364. else:
  365. raise ValueError("Unsupported backend type")
  366. def read(self, in_path):
  367. audio, audio_sample_rate = self._backend.read_file(str(in_path))
  368. return audio, audio_sample_rate
  369. class _AudioReaderBackend(_BaseReaderBackend):
  370. """_AudioReaderBackend"""
  371. class WAVReaderBackend(_AudioReaderBackend):
  372. """PandasCSVReaderBackend"""
  373. def __init__(self):
  374. super().__init__()
  375. def read_file(self, in_path):
  376. """read wav file from path"""
  377. audio, audio_sample_rate = soundfile.read(
  378. in_path, dtype="float32", always_2d=True
  379. )
  380. return audio, audio_sample_rate