From aff453a0ee663d12153e2dfad37dcd57ff716763 Mon Sep 17 00:00:00 2001 From: zhengzangw Date: Sat, 1 Jun 2024 17:43:20 +0000 Subject: [PATCH] [fix] aesthetic score error for av1 --- tools/datasets/utils.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tools/datasets/utils.py b/tools/datasets/utils.py index c91691b..fe723bf 100644 --- a/tools/datasets/utils.py +++ b/tools/datasets/utils.py @@ -93,16 +93,28 @@ def extract_frames( idx = total_frames - 1 cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + # HACK: sometimes OpenCV fails to read frames, return a black frame instead try: ret, frame = cap.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = Image.fromarray(frame) except Exception as e: - print(f"Error reading frame {video_path}: {e}") - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - frame = Image.new("RGB", (width, height), (0, 0, 0)) + print(f"[Warning] Error reading frame {idx} from {video_path}: {e}") + # First, try to read the first frame + try: + print(f"[Warning] Try reading first frame.") + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + ret, frame = cap.read() + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = Image.fromarray(frame) + # If that fails, return a black frame + except Exception as e: + print(f"[Warning] Error in reading first frame from {video_path}: {e}") + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame = Image.new("RGB", (width, height), (0, 0, 0)) + # HACK: if height or width is 0, return a black frame instead if frame.height == 0 or frame.width == 0: height = width = 256