Update app.py

app.py

@@ -14,11 +14,10 @@ LOG_FOLDER = Path("logs")
 LOG_FOLDER.mkdir(parents=True, exist_ok=True)
 TOKEN = os.environ.get("HF_TOKEN")
 
-# --- 1.5 自动下载数据 (关键修复) ---
-# 只有当本地没有数据时才尝试下载
+# --- 2. 自动下载数据 (保证不为空) ---
 if not os.path.exists(DATA_FOLDER) or not os.listdir(DATA_FOLDER):
     try:
-        print("正在从 Dataset 下载图片数据...")
+        print("🚀 正在从 Dataset 下载数据...")
         snapshot_download(
             repo_id=DATASET_REPO_ID,
             repo_type="dataset",
@@ -26,25 +25,24 @@ if not os.path.exists(DATA_FOLDER) or not os.listdir(DATA_FOLDER):
             token=TOKEN,
             allow_patterns=["*.jpg", "*.png", "*.jpeg", "*.webp", "*.txt"]
         )
-        print("下载完成！")
+        print("✅ 数据下载完成！")
     except Exception as e:
-        print(f"下载数据失败 (可能是Token问题或Repo不存在): {e}")
+        print(f"⚠️ 下载失败 (请检查 Token 或网络): {e}")
 
-# --- 2. 启动同步 ---
+# --- 3. 启动同步调度器 ---
 scheduler = CommitScheduler(
     repo_id=DATASET_REPO_ID,
     repo_type="dataset",
     folder_path=LOG_FOLDER,
-    path_in_repo="logs", # 建议把日志分文件夹存
+    path_in_repo="logs",
     every=1,
     token=TOKEN
 )
 
-# --- 3. 数据加载 ---
+# --- 4. 数据加载 ---
 def load_data():
     groups = {}
     if not os.path.exists(DATA_FOLDER):
-        os.makedirs(DATA_FOLDER, exist_ok=True)
         return {}, []
 
     for filename in os.listdir(DATA_FOLDER):
@@ -80,145 +78,180 @@ def load_data():
 
 ALL_GROUPS, ALL_GROUP_IDS = load_data()
 
-# --- 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
-    row = [user_id, datetime.now().strftime("%Y-%m-%d %H:%M:%S"), group_id, choice_labels, method_names]
-    
-    with scheduler.lock:
-        file_exists = user_file_path.exists()
-        with user_file_path.open("a", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            if not file_exists:
-                writer.writerow(["user_id", "timestamp", "group_id", "selected_labels", "selected_methods"])
-            writer.writerow(row)
-    print(f"Saved: {user_id} -> {choice_labels}")
-
 # --- 5. 核心逻辑 ---
 
-def init_state():
-    return {"user_id": str(uuid.uuid4())[:8], "index": 0, "is_finished": False}
-
-def next_question_data(user_state):
+def get_next_question(user_state):
+    """准备下一题的数据"""
     idx = user_state["index"]
+    
+    # 1. 结束判断
     if idx >= len(ALL_GROUP_IDS):
-        user_state["is_finished"] = True
-        return user_state, None, [], [] 
-
+        return (
+            gr.update(visible=False), # Origin
+            gr.update(visible=False), # Gallery
+            gr.update(visible=False), # Checkbox
+            gr.update(visible=False), # Instruction
+            gr.update(visible=False), # Submit Btn
+            gr.update(visible=False), # None Btn
+            gr.update(value="## 🎉 测试结束！感谢您的参与。", visible=True),
+            user_state,
+            [] # current_candidates_info
+        )
+    
+    # 2. 获取数据
     group_id = ALL_GROUP_IDS[idx]
     group_data = ALL_GROUPS[group_id]
     
+    # 3. 准备原图
     origin_path = group_data["origin"]
+    
+    # 4. 准备候选图（打乱）
     candidates = group_data["candidates"].copy()
     random.shuffle(candidates)
     
-    candidate_info = []
+    # 构造 Gallery 数据 [(path, label), ...] 和 选项列表 ["Option A", ...]
+    gallery_items = []
+    choices = []
+    candidates_info = [] # 用于存储真实路径，方便后续查找
+    
     for i, path in enumerate(candidates):
-        label = f"Option {chr(65+i)}"
-        candidate_info.append({"path": path, "label": label})
-
-    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, []
+        label = f"Option {chr(65+i)}" # Option A, Option B...
+        gallery_items.append((path, label))
+        choices.append(label)
+        candidates_info.append({"label": label, "path": path})
+        
+    instruction = f"### 任务 ({idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
+    
+    return (
+        gr.update(value=origin_path, visible=True if origin_path else False),
+        gr.update(value=gallery_items, visible=True),
+        gr.update(choices=choices, value=[], visible=True), # 重置多选框
+        gr.update(value=instruction, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+        user_state,
+        candidates_info
+    )
 
+def save_and_next(user_state, candidates_info, selected_options, is_none=False):
+    """保存并进入下一题"""
     current_idx = user_state["index"]
     group_id = ALL_GROUP_IDS[current_idx]
-
+    
+    # --- 保存逻辑 ---
     if is_none:
-        save_user_vote(user_state["user_id"], group_id, "Rejected All", "None_Satisfied")
+        choice_str = "Rejected All"
+        method_str = "None_Satisfied"
     else:
-        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)
-            method = parts[1] if len(parts) > 1 else name_no_ext
-            methods.append(method)
+        # 检查是否未选
+        if not selected_options:
+            raise gr.Error("请至少勾选一个选项，或点击“都不满意”")
             
-        save_user_vote(user_state["user_id"], group_id, "; ".join(labels), "; ".join(methods))
+        choice_str = "; ".join(selected_options)
+        
+        # 查找对应的方法名
+        selected_methods = []
+        for opt in selected_options: # opt is "Option A"
+            # 找到对应的文件路径
+            for info in candidates_info:
+                if info["label"] == opt:
+                    path = info["path"]
+                    filename = os.path.basename(path)
+                    name = os.path.splitext(filename)[0]
+                    # 简单清洗文件名拿到方法名
+                    parts = name.split('_', 1)
+                    method = parts[1] if len(parts) > 1 else name
+                    selected_methods.append(method)
+                    break
+        method_str = "; ".join(selected_methods)
 
+    # 写入 CSV
+    user_file = LOG_FOLDER / f"user_{user_state['user_id']}.csv"
+    with scheduler.lock:
+        exists = user_file.exists()
+        with open(user_file, "a", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            if not exists:
+                writer.writerow(["user_id", "timestamp", "group_id", "choices", "methods"])
+            writer.writerow([
+                user_state["user_id"], 
+                datetime.now().strftime("%Y-%m-%d %H:%M:%S"), 
+                group_id, 
+                choice_str, 
+                method_str
+            ])
+            
+    # --- 状态更新 ---
     user_state["index"] += 1
-    return user_state, []
+    
+    # --- 加载下一题 ---
+    return get_next_question(user_state)
 
-# --- 6. 界面构建 (Gradio 5.0+) ---
+# --- 6. 界面构建 ---
 with gr.Blocks(title="User Study") as demo:
     
-    state_main = gr.State(init_state) # 修正：使用 init_state 函数引用
-    state_origin = gr.State()
-    state_candidates = gr.State([])
-    state_instruction = gr.State("")
-    state_selection = gr.State([])
-
-    # @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):
-        if main_st["is_finished"]:
-            gr.Markdown("## 🎉 测试结束！\n感谢您的参与，所有结果已保存。")
-            return
-
-        idx = main_st["index"]
-        total = len(ALL_GROUP_IDS)
-        gr.Markdown(f"### 任务 ({idx + 1} / {total if total > 0 else 0})\n\n{instruction}")
+    # 状态变量
+    state_user = gr.State({"user_id": str(uuid.uuid4())[:8], "index": 0})
+    state_candidates_info = gr.State([]) # 存当前页面的候选图信息，用于把 Option A 映射回文件名
+    
+    # 布局
+    with gr.Row():
+        md_instruction = gr.Markdown("Loading...")
+    
+    with gr.Row():
+        # 左侧：原图
+        with gr.Column(scale=1):
+            img_origin = gr.Image(label="Reference (参考原图)", interactive=False, height=400)
         
-        if total == 0:
-            gr.Markdown("⚠️ **错误**：未找到任何数据。请检查 Dataset 设置。")
-            return
-
-        if origin:
+        # 右侧：候选图 + 选择区
+        with gr.Column(scale=2):
+            # 1. Gallery 显示所有选项
+            gallery_candidates = gr.Gallery(
+                label="Candidates (候选结果)", 
+                columns=[2], 
+                height="auto",
+                object_fit="contain",
+                interactive=False # 禁止 Gallery 自身的点击选择，用下面的 Checkbox 替代
+            )
+            
+            gr.Markdown("👇 **请在下方勾选您认为最好的结果（可多选）：**")
+            
+            # 2. 多选框 (核心交互)
+            checkbox_options = gr.CheckboxGroup(
+                choices=[], 
+                label="您的选择", 
+                info="对应上方图片的标签 (Option A, B...)"
+            )
+            
             with gr.Row():
-                with gr.Column(scale=1):
-                    gr.Image(origin, label="Reference (参考原图)", interactive=False, height=300)
-                with gr.Column(scale=2):
-                    gr.Markdown("👈 **请参考左侧原图**，并在下方选择您认为质量最好的图片（可多选）。")
-
-        with gr.Row(wrap=True):
-            for i, item in enumerate(candidates):
-                is_selected = i in selection
-                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)
-                        current_sel.sort()
-                        return current_sel
-                    
-                    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)
-
-    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])
-
-    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(on_submit, inputs=[state_main, state_candidates, state_selection], outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction])
+                btn_submit = gr.Button("🚀 提交 (Submit)", variant="primary")
+                btn_none = gr.Button("🚫 都不满意 (None)", variant="stop")
+    
+    md_end = gr.Markdown(visible=False)
 
-    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(on_none, inputs=[state_main, state_candidates, state_selection], outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction])
+    # --- 事件绑定 ---
+    
+    # 1. 页面加载时，加载第一题
+    demo.load(
+        fn=get_next_question,
+        inputs=[state_user],
+        outputs=[img_origin, gallery_candidates, checkbox_options, md_instruction, btn_submit, btn_none, md_end, state_user, state_candidates_info]
+    )
+    
+    # 2. 提交按钮
+    btn_submit.click(
+        fn=lambda s, c, o: save_and_next(s, c, o, is_none=False),
+        inputs=[state_user, state_candidates_info, checkbox_options],
+        outputs=[img_origin, gallery_candidates, checkbox_options, md_instruction, btn_submit, btn_none, md_end, state_user, state_candidates_info]
+    )
+    
+    # 3. 都不满意按钮
+    btn_none.click(
+        fn=lambda s, c, o: save_and_next(s, c, o, is_none=True),
+        inputs=[state_user, state_candidates_info, checkbox_options],
+        outputs=[img_origin, gallery_candidates, checkbox_options, md_instruction, btn_submit, btn_none, md_end, state_user, state_candidates_info]
+    )
 
 if __name__ == "__main__":
     demo.launch()