Update app.py

app.py

@@ -54,7 +54,12 @@ with open("videos_by_dimension.json") as f:
 types = ['appearance_style', 'color', 'temporal_style', 'spatial_relationship', 'temporal_flickering', 'scene', 'multiple_objects', 'object_class', 'human_action', 'overall_consistency', 'subject_consistency']
 
 def get_video_path_local(model_name, type, prompt):
-    video_path_subfolder = os.path.join(model_name, type)
+    if 'Show-1' in model_name:
+        video_path_subfolder = os.path.join(model_name, type, 'super2')
+    elif 'videocrafter-1' in model_name:
+        video_path_subfolder = os.path.join(model_name, type, '1024x576')
+    else:
+        video_path_subfolder = os.path.join(model_name, type)
     try:
         return hf_api.hf_hub_download(
             repo_id = repo_id,
@@ -108,14 +113,28 @@ def display_videos(type, prompt, model_name_1, model_name_2):
     video_path2 = get_video_path_local(model_name_2, type, prompt)
     return video_path1, video_path2
 
+def record_user_feedback_a(model_name1, model_name2, type, prompt):
+    # 0 means model A  better, 1 means model B better 
+    with open("user_feedback.csv",'a') as f:
+        f.write(f"{model_name1}\t{model_name2}\t{type}\t{prompt}\t{0}\n")
+def record_user_feedback_b(model_name1, model_name2, type, prompt):
+    # 0 means model A  better, 1 means model B better 
+    with open("user_feedback.csv",'a') as f:
+        f.write(f"{model_name1}\t{model_name2}\t{type}\t{prompt}\t{1}\n")  
 
 with gr.Blocks() as interface:
+    gr.Markdown("# VBench Video Arena")
+    gr.Markdown("""
+**Random 2 videos** for randomly picking two models to compare random with the same dimension and prompt
+**Play Selection** is used for the user to select the model, dimension, prompt in the drop-down box, and display them 
+If you are interested, you can also leave your comments.""")
+    
     type_output = gr.Dropdown(label="Type", choices=types, value=types[0])
     prompt_output = gr.Dropdown(label="Prompt", choices=dimension[types[0]], value=dimension[types[0]][0])
     prompt_placeholder = gr.State()
     with gr.Row():
         random_button = gr.Button("🎲 Random 2 videos")
-        display_button = gr.Button("⏸ Play options")
+        display_button = gr.Button("🎇 Play Selection")
    
     with gr.Row():
         with gr.Column():
@@ -124,9 +143,11 @@ with gr.Blocks() as interface:
         with gr.Column():
             model_name_2_output = gr.Dropdown(label="Model Name 2", choices=model_names, value=model_names[1])
             video_output_2 = gr.Video(label="Video 2")
-    
+    with gr.Row():
+        feed0 = gr.Button("👈 Model A is better")
+        feed1 = gr.Button("👉 Model B is better")
     type_output.change(fn=update_prompt_options, inputs=[type_output], outputs=[prompt_output])
-
+    
     
     random_button.click(
         fn=get_random_video,
@@ -144,4 +165,13 @@ with gr.Blocks() as interface:
         outputs=[video_output_1, video_output_2]
     )
 
+    feed0.click(
+        fn = record_user_feedback_a,
+        inputs=[model_name_1_output, model_name_2_output, type_output, prompt_placeholder]
+    )
+    feed1.click(
+        fn = record_user_feedback_b,
+        inputs=[model_name_1_output, model_name_2_output, type_output, prompt_placeholder]
+    )
+
 interface.launch()