diff --git a/eval/vbench_i2v/calc_vbench_i2v.py b/eval/vbench_i2v/calc_vbench_i2v.py index a96c812..4bc2297 100644 --- a/eval/vbench_i2v/calc_vbench_i2v.py +++ b/eval/vbench_i2v/calc_vbench_i2v.py @@ -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)