72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
import argparse
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
|
|
from vbench import VBench
|
|
from vbench2_beta_i2v import VBenchI2V
|
|
|
|
full_info_path = "eval/vbench_i2v/vbench2_i2v_full_info.json"
|
|
video_quality_dimensions = [
|
|
"subject_consistency",
|
|
"background_consistency",
|
|
"motion_smoothness",
|
|
"dynamic_degree",
|
|
"aesthetic_quality",
|
|
"imaging_quality",
|
|
"temporal_flickering",
|
|
]
|
|
i2v_dimensions = ["i2v_subject", "i2v_background", "camera_motion"]
|
|
|
|
|
|
def str2bool(v):
|
|
if isinstance(v, bool):
|
|
return v
|
|
if v.lower() in ("yes", "true", "t", "y", "1"):
|
|
return True
|
|
elif v.lower() in ("no", "false", "f", "n", "0"):
|
|
return False
|
|
else:
|
|
raise argparse.ArgumentTypeError("Boolean value expected.")
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("video_folder", type=str) # samples/samples..._vbench_i2v/
|
|
parser.add_argument("model_ckpt", type=str)
|
|
parser.add_argument("--start", type=int, default=0) # start index of dimension to be evaluated
|
|
parser.add_argument("--end", type=int, default=-1) # start index of dimension to be evaluated
|
|
parser.add_argument("--calc_i2v", type=str2bool, default=True)
|
|
parser.add_argument("--calc_quality", type=str2bool, default=True)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
output_dir = os.path.join(args.model_ckpt, "vbench_i2v")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
video_path = args.video_folder
|
|
|
|
start_time = time.time()
|
|
|
|
if args.calc_i2v:
|
|
my_VBench_I2V = VBenchI2V(torch.device("cuda"), full_info_path, output_dir)
|
|
end = args.end if args.end != -1 else len(i2v_dimensions)
|
|
for i2v_dim in i2v_dimensions[args.start : end]:
|
|
my_VBench_I2V.evaluate(videos_path=video_path, name=i2v_dim, dimension_list=[i2v_dim], resolution="1-1")
|
|
|
|
kwargs = {}
|
|
kwargs["imaging_quality_preprocessing_mode"] = "longer" # use VBench/evaluate.py default
|
|
|
|
if args.calc_quality:
|
|
my_VBench = VBench(torch.device("cuda"), full_info_path, output_dir)
|
|
end = args.end if args.end != -1 else len(video_quality_dimensions)
|
|
for quality_dim in video_quality_dimensions[args.start : end]:
|
|
my_VBench.evaluate(
|
|
videos_path=video_path, name=quality_dim, dimension_list=[quality_dim], mode="vbench_standard", **kwargs
|
|
)
|
|
|
|
print("Runtime: %s seconds " % (time.time() - start_time))
|