nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-15 11:52:47 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[fix] aesthetic score error for av1

Browse source
This commit is contained in:
zhengzangw 2024-06-01 17:43:20 +00:00
parent 735a964743
commit aff453a0ee
1 changed files with 16 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
tools/datasets/utils.py
View file

@ -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