nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-10 21:01:26 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[fix] read_video for tools

Browse source
This commit is contained in:
zhengzangw 2024-05-31 07:17:01 +00:00
parent fd51586704
commit 2988d5be99
2 changed files with 5 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
opensora/datasets/read_video.py
View file

@ -172,7 +172,7 @@ def read_video_cv2(video_path):
return frames, vinfo
def read_video(video_path, backend='cv2'):
def read_video(video_path, backend='av'):
if backend == 'cv2':
vframes, vinfo = read_video_cv2(video_path)
elif backend == 'av':

6
tools/datasets/datautil.py
View file

@ -10,11 +10,12 @@ from glob import glob
import cv2
import numpy as np
import pandas as pd
import torchvision
from PIL import Image
from tqdm import tqdm
import torchvision
from .utils import IMG_EXTENSIONS
from opensora.datasets.read_video import read_video
tqdm.pandas()
@ -116,7 +117,8 @@ def get_image_info(path, backend="pillow"):
def get_video_info(path, backend="torchvision"):
if backend == "torchvision":
try:
vframes, _, infos = torchvision.io.read_video(filename=path, pts_unit="sec", output_format="TCHW")
# vframes, infos = read_video(path)
vframes, _, infos = torchvision.io.read_video(filename=path, pts_unit="sec", output_format="THWC")
num_frames, height, width = vframes.shape[0], vframes.shape[2], vframes.shape[3]
if "video_fps" in infos:
fps = infos["video_fps"]