Update app.py

app.py

@@ -5,7 +5,7 @@ import uuid
 import csv
 from datetime import datetime
 from pathlib import Path
-from huggingface_hub import CommitScheduler
+from huggingface_hub import CommitScheduler, snapshot_download
 
 # --- 1. 配置区域 ---
 DATASET_REPO_ID = "Emilyxml/moveit"
@@ -14,17 +14,33 @@ LOG_FOLDER = Path("logs")
 LOG_FOLDER.mkdir(parents=True, exist_ok=True)
 TOKEN = os.environ.get("HF_TOKEN")
 
+# --- 1.5 自动下载数据 (关键修复) ---
+# 只有当本地没有数据时才尝试下载
+if not os.path.exists(DATA_FOLDER) or not os.listdir(DATA_FOLDER):
+    try:
+        print("正在从 Dataset 下载图片数据...")
+        snapshot_download(
+            repo_id=DATASET_REPO_ID,
+            repo_type="dataset",
+            local_dir=DATA_FOLDER,
+            token=TOKEN,
+            allow_patterns=["*.jpg", "*.png", "*.jpeg", "*.webp", "*.txt"]
+        )
+        print("下载完成！")
+    except Exception as e:
+        print(f"下载数据失败 (可能是Token问题或Repo不存在): {e}")
+
 # --- 2. 启动同步 ---
 scheduler = CommitScheduler(
     repo_id=DATASET_REPO_ID,
     repo_type="dataset",
     folder_path=LOG_FOLDER,
-    path_in_repo="data",
+    path_in_repo="logs", # 建议把日志分文件夹存
     every=1,
     token=TOKEN
 )
 
-# --- 3. 数据加载 (不变) ---
+# --- 3. 数据加载 ---
 def load_data():
     groups = {}
     if not os.path.exists(DATA_FOLDER):
@@ -64,7 +80,7 @@ def load_data():
 
 ALL_GROUPS, ALL_GROUP_IDS = load_data()
 
-# --- 4. 保存逻辑 (不变) ---
+# --- 4. 保存逻辑 ---
 def save_user_vote(user_id, group_id, choice_labels, method_names):
     user_filename = f"user_{user_id}.csv"
     user_file_path = LOG_FOLDER / user_filename
@@ -79,38 +95,24 @@ def save_user_vote(user_id, group_id, choice_labels, method_names):
             writer.writerow(row)
     print(f"Saved: {user_id} -> {choice_labels}")
 
-# --- 5. 核心逻辑：状态管理 ---
+# --- 5. 核心逻辑 ---
 
 def init_state():
-    """初始化用户状态"""
-    return {
-        "user_id": str(uuid.uuid4())[:8],
-        "index": 0,
-        "is_finished": False
-    }
+    return {"user_id": str(uuid.uuid4())[:8], "index": 0, "is_finished": False}
 
 def next_question_data(user_state):
-    """
-    计算下一题的数据，并返回给前端渲染。
-    这里只处理数据准备（打乱顺序等），不处理 UI。
-    """
     idx = user_state["index"]
-    
     if idx >= len(ALL_GROUP_IDS):
         user_state["is_finished"] = True
-        return user_state, None, [], [] # 结束状态
+        return user_state, None, [], [] 
 
     group_id = ALL_GROUP_IDS[idx]
     group_data = ALL_GROUPS[group_id]
     
-    # 准备原图
     origin_path = group_data["origin"]
-    
-    # 准备候选图（打乱）
     candidates = group_data["candidates"].copy()
     random.shuffle(candidates)
     
-    # 构造候选图信息列表 [(path, "Option A"), (path, "Option B")...]
     candidate_info = []
     for i, path in enumerate(candidates):
         label = f"Option {chr(65+i)}"
@@ -119,28 +121,22 @@ def next_question_data(user_state):
     return user_state, origin_path, candidate_info, group_data["instruction"]
 
 def submit_logic(user_state, current_candidates, selected_indices, is_none=False):
-    """处理提交，保存数据，并进入下一题"""
     if user_state["is_finished"]:
-        return user_state, [], []
+        return user_state, []
 
     current_idx = user_state["index"]
     group_id = ALL_GROUP_IDS[current_idx]
 
-    # 1. 保存数据
     if is_none:
         save_user_vote(user_state["user_id"], group_id, "Rejected All", "None_Satisfied")
     else:
-        # 如果不是 None 但没选，直接返回不跳转（或者你可以允许）
         if not selected_indices:
             raise gr.Error("请至少选择一张图片，或点击“都不满意”")
-            
         labels = []
         methods = []
         for idx in selected_indices:
             info = current_candidates[idx]
             labels.append(info["label"])
-            
-            # 提取方法名
             filename = os.path.basename(info["path"])
             name_no_ext = os.path.splitext(filename)[0]
             parts = name_no_ext.split('_', 1)
@@ -149,115 +145,80 @@ def submit_logic(user_state, current_candidates, selected_indices, is_none=False
             
         save_user_vote(user_state["user_id"], group_id, "; ".join(labels), "; ".join(methods))
 
-    # 2. 索引+1
     user_state["index"] += 1
-    
-    # 3. 清空选择
-    return user_state, [] # 返回新的 state 和空的 selected_indices
+    return user_state, []
 
-# --- 6. 界面构建 (使用 @gr.render 实现自定义布局) ---
+# --- 6. 界面构建 (Gradio 5.0+) ---
 with gr.Blocks(title="User Study") as demo:
     
-    # === 状态变量 ===
-    # state_main: 存 user_id, index
-    state_main = gr.State(init_state())
-    # state_current_data: 存当前题目的数据 (原图, 候选图列表, 说明文)
+    state_main = gr.State(init_state) # 修正：使用 init_state 函数引用
     state_origin = gr.State()
-    state_candidates = gr.State([]) # list of dict
+    state_candidates = gr.State([])
     state_instruction = gr.State("")
-    # state_selection: 存当前选中的索引 [0, 2]
     state_selection = gr.State([])
 
-    # === 页面布局 ===
-    header = gr.Markdown("Loading...")
-    
-    # 动态渲染区域
+    # @gr.render 需要 Gradio 5.0+
     @gr.render(inputs=[state_main, state_origin, state_candidates, state_selection, state_instruction])
     def render_content(main_st, origin, candidates, selection, instruction):
-        
-        # 1. 如果结束了
         if main_st["is_finished"]:
             gr.Markdown("## 🎉 测试结束！\n感谢您的参与，所有结果已保存。")
             return
 
-        # 2. 显示说明
         idx = main_st["index"]
-        gr.Markdown(f"### 任务 ({idx + 1} / {len(ALL_GROUP_IDS)})\n\n{instruction}")
+        total = len(ALL_GROUP_IDS)
+        gr.Markdown(f"### 任务 ({idx + 1} / {total if total > 0 else 0})\n\n{instruction}")
         
-        # 3. 显示原图 (Reference)
+        if total == 0:
+            gr.Markdown("⚠️ **错误**：未找到任何数据。请检查 Dataset 设置。")
+            return
+
         if origin:
             with gr.Row():
                 with gr.Column(scale=1):
                     gr.Image(origin, label="Reference (参考原图)", interactive=False, height=300)
                 with gr.Column(scale=2):
-                    gr.Markdown("👈 **请参考左侧原图**，并在下方选择您认为质量最好的图片（可多选）。\n\n**点击图片下方的按钮进行选择。**")
+                    gr.Markdown("👈 **请参考左侧原图**，并在下方选择您认为质量最好的图片（可多选）。")
 
-        # 4. 显示候选图 (Grid Layout)
-        # 使用 Row wrap=True 实现自动换行
         with gr.Row(wrap=True):
             for i, item in enumerate(candidates):
                 is_selected = i in selection
-                
-                # 定义每个卡片的样式
-                with gr.Column(min_width=200): # 限制最小宽度，类似 Gallery 效果
+                with gr.Column(min_width=200): 
                     gr.Image(item["path"], show_label=False, interactive=False)
-                    
-                    # === 核心：每个图片下的按钮 ===
                     btn_text = f"✅ {item['label']} (已选)" if is_selected else f"⬜️ {item['label']} (点击选择)"
                     btn_variant = "primary" if is_selected else "secondary"
                     
                     btn = gr.Button(btn_text, variant=btn_variant)
                     
-                    # 按钮点击逻辑：切换选中状态
                     def toggle(idx, current_sel):
-                        if idx in current_sel:
-                            current_sel.remove(idx)
-                        else:
-                            current_sel.append(idx)
+                        if idx in current_sel: current_sel.remove(idx)
+                        else: current_sel.append(idx)
                         current_sel.sort()
                         return current_sel
                     
-                    # 绑定点击事件，更新 state_selection，从而触发重绘
                     btn.click(fn=toggle, inputs=[gr.Number(i, visible=False), state_selection], outputs=[state_selection])
 
-    # === 底部操作栏 ===
     with gr.Row():
         btn_submit = gr.Button("🚀 提交 (Submit)", variant="primary", scale=2)
         btn_none = gr.Button("🚫 都不满意 (None)", variant="stop", scale=1)
 
-    # === 事件流 ===
-    
-    # 1. 初始化加载第一题
     def load_first(main_st):
         return next_question_data(main_st)
 
     demo.load(load_first, inputs=[state_main], outputs=[state_main, state_origin, state_candidates, state_instruction])
 
-    # 2. 提交按钮 -> 保存 -> 准备下一题 -> 清空选择
     def on_submit(main_st, cands, sel):
-        # 先保存
         new_main, new_sel = submit_logic(main_st, cands, sel, is_none=False)
-        # 再加载下一题数据
         updated_main, origin, new_cands, instr = next_question_data(new_main)
         return updated_main, new_sel, origin, new_cands, instr
 
-    btn_submit.click(
-        fn=on_submit,
-        inputs=[state_main, state_candidates, state_selection],
-        outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction]
-    )
+    btn_submit.click(on_submit, inputs=[state_main, state_candidates, state_selection], outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction])
 
-    # 3. 都不满意 -> 保存 -> 准备下一题
     def on_none(main_st, cands, sel):
         new_main, new_sel = submit_logic(main_st, cands, sel, is_none=True)
         updated_main, origin, new_cands, instr = next_question_data(new_main)
         return updated_main, new_sel, origin, new_cands, instr
         
-    btn_none.click(
-        fn=on_none,
-        inputs=[state_main, state_candidates, state_selection],
-        outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction]
-    )
+    btn_none.click(on_none, inputs=[state_main, state_candidates, state_selection], outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction])
 
 if __name__ == "__main__":
     demo.launch()