Update app.py

app.py

@@ -2,17 +2,47 @@ import gradio as gr
 import random
 from datasets import load_dataset
 import os
-hf_token = os.environ['hf_token']
 
-# 加载数据集
-dataset = load_dataset("Vchitect/VBench_sampled_video")
-# 随机选择一个视频
+hf_token = os.environ['hf_token']  # 确保环境变量中有你的令牌
+submission_url = "Vchitect/VBench_sampled_video"  # 数据集的 URL
+local_dir = "VBench_sampled_video"  # 本地文件夹路径
+
+# 克隆数据集
+submission_repo = Repository(local_dir=local_dir, clone_from=submission_url, use_auth_token=hf_token, repo_type="dataset")
+submission_repo.git_pull()  # 更新本地仓库
+
+model_names = os.listdir(local_dir)
+
+with open("videos_by_dimension.json") as f:
+    dimension = json.load(f)['videos_by_dimension']
+
+# with open("all_videos.json") as f:
+    # all_videos = json.load(f)
+
+types = ['appearance_style', 'color', 'temporal_style', 'spatial_relationship', 'temporal_flickering', 'scene', 'multiple_objects', 'object_class', 'human_action', 'overall_consistency', 'subject_consistency']
+
 def get_random_video():
     # 随机选择一个索引
-    random_index = random.randint(0, len(dataset['train']) - 1)
-    # 获取视频路径
-    print(dataset['train'][random_index])
-    video_path = dataset['train'][random_index]['video_path']
+    random_index = random.randint(0, len(types) - 1)
+    type = types[random_index]
+    # 随机选择一个Prompt
+    random_index = random.randint(0, len(dimension[type]) - 1)
+    prompt = dimension[type][random_index]
+    # 随机一个模型
+    random_index = random.randint(0, len(model_names) - 1)
+    model_name =  model_names[random_index]
+
+    video_path = os.path.join(model_name, type, prompt)
+    if os.path.exists(video_path):
+        print(video_path)
+        return video_path
+    else:
+        video_path = os.path.join(model_name, prompt)
+        if os.path.exists(video_path):
+            print(video_path)
+            return video_path
+    # video_path = dataset['train'][random_index]['video_path']
+    print('error:', video_path)
     return video_path
 
 # Gradio 接口