From 67f03c6cdaf2ee2889eaabf5a02e222452acb44d Mon Sep 17 00:00:00 2001
From: zhengzangw <zhengzangw@gmail.com>
Date: Sat, 8 Jun 2024 15:46:42 +0000
Subject: [PATCH] update vbench script

---
 .gitignore                 |  1 +
 eval/loss/launch.sh        |  3 --
 eval/vbench/calc_vbench.py | 66 ++++++++++++++++++++++++++++++++++++++
 eval/vbench/launch.sh      |  3 --
 eval/vbench/launch_calc.sh | 17 ++++++++++
 5 files changed, 84 insertions(+), 6 deletions(-)
 create mode 100644 eval/vbench/calc_vbench.py
 create mode 100644 eval/vbench/launch_calc.sh

diff --git a/.gitignore b/.gitignore
index b4237e0..57b6f55 100644
--- a/.gitignore
+++ b/.gitignore
@@ -199,4 +199,5 @@ tools/caption/pllava_dir/PLLaVA/
 
 # vbench
 vbench
+!eval/vbench
 vbench2_beta_i2v
diff --git a/eval/loss/launch.sh b/eval/loss/launch.sh
index a80f791..67d230a 100644
--- a/eval/loss/launch.sh
+++ b/eval/loss/launch.sh
@@ -1,8 +1,5 @@
 #!/bin/bash
 
-set -x
-set -e
-
 CMD="torchrun --standalone --nproc_per_node 1 eval/loss/eval_loss.py configs/opensora-v1-2/misc/eval_loss.py"
 CKPT_PATH=$1
 MODEL_NAME=$2
diff --git a/eval/vbench/calc_vbench.py b/eval/vbench/calc_vbench.py
new file mode 100644
index 0000000..b2505bf
--- /dev/null
+++ b/eval/vbench/calc_vbench.py
@@ -0,0 +1,66 @@
+import argparse
+import os
+from vbench import VBench
+import torch
+import time
+
+full_info_path = "eval/vbench/VBench_full_info.json"
+dimensions = [
+    # Quality Score
+    "subject_consistency",
+    "background_consistency",
+    "motion_smoothness",
+    "dynamic_degree",
+    "aesthetic_quality",
+    "imaging_quality",
+    "temporal_flickering",
+    # Semantic Score
+    "object_class",
+    "multiple_objects",
+    "color",
+    "spatial_relationship",
+    "scene",
+    "temporal_style",
+    "overall_consistency",
+    "human_action",
+    "appearance_style",
+]
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("video_folder", type=str)  # samples/samples..._vbench/eval
+    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
+
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    output_dir = os.path.join(args.model_ckpt, "vbench")
+    os.makedirs(output_dir, exist_ok=True)
+    video_path = args.video_folder
+
+    kwargs = {}
+    kwargs['imaging_quality_preprocessing_mode'] = 'longer' # use VBench/evaluate.py default
+
+    start_time = time.time()
+
+    # NOTE: important to use torch.device("cuda"), else will have issue with object_class third_party module
+    my_VBench = VBench(torch.device("cuda"), full_info_path, output_dir)
+    if args.end == -1: # adjust end accordingly
+        args.end = len(dimensions)
+    for dim in dimensions[args.start:args.end]:
+        my_VBench.evaluate(
+            videos_path=video_path,
+            name=dim,
+            local=False,
+            read_frame=False,
+            dimension_list=[dim],
+            mode='vbench_standard',
+            **kwargs
+        )
+
+    print("Runtime: %s seconds " % (time.time() - start_time))
\ No newline at end of file
diff --git a/eval/vbench/launch.sh b/eval/vbench/launch.sh
index d45c9fb..0cb342e 100644
--- a/eval/vbench/launch.sh
+++ b/eval/vbench/launch.sh
@@ -1,8 +1,5 @@
 # !/bin/bash
 
-set -x
-set -e
-
 CKPT=$1
 NUM_FRAMES=$2
 MODEL_NAME=$3
diff --git a/eval/vbench/launch_calc.sh b/eval/vbench/launch_calc.sh
new file mode 100644
index 0000000..53114b9
--- /dev/null
+++ b/eval/vbench/launch_calc.sh
@@ -0,0 +1,17 @@
+# !/bin/bash
+
+VIDEO_DIR=$1
+CKPT_DIR=$2
+LOG_BASE=$CKPT_DIR
+mkdir -p $LOG_BASE
+echo "Logging to $LOG_BASE"
+
+GPUS=(0 1 2 3 4 5 6 7)
+START_INDEX_LIST=(0 2 4 6 8 10 12 14)
+END_INDEX_LIST=(2 4 6 8 10 12 14 16)
+TASK_ID_LIST=(calc_vbench_a calc_vbench_b calc_vbench_c calc_vbench_d calc_vbench_e calc_vbench_f calc_vbench_g calc_vbench_h) # for log records only
+
+
+for i in "${!GPUS[@]}"; do
+    CUDA_VISIBLE_DEVICES=${GPUS[i]} python eval/vbench/calc_vbench.py $VIDEO_DIR $CKPT_DIR --start ${START_INDEX_LIST[i]} --end ${END_INDEX_LIST[i]} > ${LOG_BASE}/${TASK_ID_LIST[i]}.log 2>&1 &
+done