This commit is contained in:
Shen-Chenhui 2024-06-09 02:28:08 +00:00
parent 05764ee113
commit e4d2985a64

View file

@ -20,14 +20,25 @@ video_quality_dimensions = [
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=bool, default=True)
parser.add_argument("--calc_quality", type=bool, default=True)
parser.add_argument("--calc_i2v", type=str2bool, default=True)
parser.add_argument("--calc_quality", type=str2bool, default=True)
args = parser.parse_args()
return args
@ -40,7 +51,7 @@ if __name__ == "__main__":
start_time = time.time()
if args.calc_i2v:
if args.calc_i2v == 1:
my_VBench_I2V = VBenchI2V(torch.device("cuda"), full_info_path, output_dir)
if args.end == -1: # adjust end accordingly
args.end = len(i2v_dimensions)
@ -50,7 +61,7 @@ if __name__ == "__main__":
kwargs = {}
kwargs["imaging_quality_preprocessing_mode"] = "longer" # use VBench/evaluate.py default
if args.calc_quality:
if args.calc_quality == 1:
if args.end == -1: # adjust end accordingly
args.end = len(video_quality_dimensions)
my_VBench = VBench(torch.device("cuda"), full_info_path, output_dir)