| 184 | return mask |
| 185 | |
| 186 | def read_video_frames(video_path, use_type='cv2', is_rgb=True, info=False): |
| 187 | frames = [] |
| 188 | if use_type == "decord": |
| 189 | import decord |
| 190 | decord.bridge.set_bridge("native") |
| 191 | try: |
| 192 | cap = decord.VideoReader(video_path) |
| 193 | total_frames = len(cap) |
| 194 | fps = cap.get_avg_fps() |
| 195 | height, width, _ = cap[0].shape |
| 196 | frames = [cap[i].asnumpy() for i in range(len(cap))] |
| 197 | except Exception as e: |
| 198 | print(f"Decord read error: {e}") |
| 199 | return None |
| 200 | elif use_type == "cv2": |
| 201 | try: |
| 202 | cap = cv2.VideoCapture(video_path) |
| 203 | fps = cap.get(cv2.CAP_PROP_FPS) |
| 204 | width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 205 | height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 206 | while cap.isOpened(): |
| 207 | ret, frame = cap.read() |
| 208 | if not ret: |
| 209 | break |
| 210 | if is_rgb: |
| 211 | frames.append(frame[..., ::-1]) |
| 212 | else: |
| 213 | frames.append(frame) |
| 214 | cap.release() |
| 215 | total_frames = len(frames) |
| 216 | except Exception as e: |
| 217 | print(f"OpenCV read error: {e}") |
| 218 | return None |
| 219 | else: |
| 220 | raise ValueError(f"Unknown video type {use_type}") |
| 221 | if info: |
| 222 | return frames, fps, width, height, total_frames |
| 223 | else: |
| 224 | return frames |
| 225 | |
| 226 | |
| 227 | |