diff --git a/inference_propainter.py b/inference_propainter.py index 4d7f92f..b093954 100644 --- a/inference_propainter.py +++ b/inference_propainter.py @@ -49,10 +49,16 @@ def resize_frames(frames, size=None): def read_frame_from_videos(frame_root): if frame_root.endswith(('mp4', 'mov', 'avi', 'MP4', 'MOV', 'AVI')): # input video path video_name = os.path.basename(frame_root)[:-4] - vframes, aframes, info = torchvision.io.read_video(filename=frame_root, pts_unit='sec') # RGB - frames = list(vframes.numpy()) - frames = [Image.fromarray(f) for f in frames] - fps = info['video_fps'] + # torchvision.io.read_video was removed in torchvision >= 0.23 — read via cv2 + cap = cv2.VideoCapture(frame_root) + fps = cap.get(cv2.CAP_PROP_FPS) or 24 + frames = [] + while True: + ok, frame = cap.read() + if not ok: + break + frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) + cap.release() else: video_name = os.path.basename(frame_root) frames = []