mirror of
https://github.com/hpcaitech/Open-Sora.git
synced 2026-04-10 12:49:38 +02:00
febf3ad4b2
* upload v2.0 * update docs * [hotfix] fit latest fa3 (#802) * update readme * update readme * update readme * update train readme * update readme * update readme: motion score * cleaning video dc ae WIP * update config * add dependency functions * undo cleaning * use latest dcae * complete high compression training * update hcae config * cleaned up vae * update ae.md * further cleanup * update vae & ae paths * align naming of ae * [hotfix] fix ring attn bwd for fa3 (#803) * train ae default without wandb * update config * update evaluation results * added hcae report * update readme * update readme demo * update readme demo * update readme gif * display demo directly in readme * update paper * delete files --------- Co-authored-by: Hongxin Liu <lhx0217@gmail.com> Co-authored-by: Shen-Chenhui <shen_chenhui@u.nus.edu> Co-authored-by: wuxiwen <wuxiwen.simon@gmail.com>
71 lines
1.9 KiB
Python
71 lines
1.9 KiB
Python
import argparse
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pandarallel import pandarallel
|
|
from torchvision.io.video import read_video
|
|
from tqdm import tqdm
|
|
|
|
|
|
def set_parallel(num_workers: int = None) -> callable:
|
|
if num_workers == 0:
|
|
return lambda x, *args, **kwargs: x.progress_apply(*args, **kwargs)
|
|
else:
|
|
if num_workers is not None:
|
|
pandarallel.initialize(progress_bar=True, nb_workers=num_workers)
|
|
else:
|
|
pandarallel.initialize(progress_bar=True)
|
|
return lambda x, *args, **kwargs: x.parallel_apply(*args, **kwargs)
|
|
|
|
|
|
def get_video_info(path: str) -> pd.Series:
|
|
vframes, _, vinfo = read_video(path, pts_unit="sec", output_format="TCHW")
|
|
num_frames, C, height, width = vframes.shape
|
|
fps = round(vinfo["video_fps"], 3)
|
|
aspect_ratio = height / width if width > 0 else np.nan
|
|
resolution = height * width
|
|
|
|
ret = pd.Series(
|
|
[height, width, fps, num_frames, aspect_ratio, resolution],
|
|
index=[
|
|
"height",
|
|
"width",
|
|
"fps",
|
|
"num_frames",
|
|
"aspect_ratio",
|
|
"resolution",
|
|
],
|
|
dtype=object,
|
|
)
|
|
return ret
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", type=str, required=True, help="Input file path")
|
|
parser.add_argument("--output", type=str, required=True, help="Output file path")
|
|
parser.add_argument(
|
|
"--num_workers", type=int, default=None, help="Number of workers"
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
input_path = args.input
|
|
output_path = args.output
|
|
num_workers = args.num_workers
|
|
|
|
df = pd.read_csv(input_path)
|
|
tqdm.pandas()
|
|
apply = set_parallel(num_workers)
|
|
|
|
result = apply(df["path"], get_video_info)
|
|
for col in result.columns:
|
|
df[col] = result[col]
|
|
df.to_csv(output_path, index=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|