Update app.py

app.py

@@ -8,13 +8,13 @@ from pathlib import Path
 from huggingface_hub import CommitScheduler
 
 # --- 1. 配置区域 ---
-DATASET_REPO_ID = "Emilyxml/moveit"  # 你的数据集
-DATA_FOLDER = "data"                 # 数据文件夹
-LOG_FOLDER = Path("logs")            # 本地日志
+DATASET_REPO_ID = "Emilyxml/moveit"
+DATA_FOLDER = "data"
+LOG_FOLDER = Path("logs")
 LOG_FOLDER.mkdir(parents=True, exist_ok=True)
 TOKEN = os.environ.get("HF_TOKEN")
 
-# --- 2. 启动同步调度器 ---
+# --- 2. 启动同步 ---
 scheduler = CommitScheduler(
     repo_id=DATASET_REPO_ID,
     repo_type="dataset",
@@ -24,7 +24,7 @@ scheduler = CommitScheduler(
     token=TOKEN
 )
 
-# --- 3. 数据加载逻辑 (区分原图和候选图) ---
+# --- 3. 数据加载 (不变) ---
 def load_data():
     groups = {}
     if not os.path.exists(DATA_FOLDER):
@@ -34,21 +34,16 @@ def load_data():
     for filename in os.listdir(DATA_FOLDER):
         if filename.startswith('.'): continue
         file_path = os.path.join(DATA_FOLDER, filename)
-        prefix = filename[:5] # 以前5个字符作为组ID
+        prefix = filename[:5]
         
         if prefix not in groups:
-            # origin 存原图路径，candidates 存其他方法的图
             groups[prefix] = {"origin": None, "candidates": [], "instruction": "暂无说明"}
 
-        # 图片处理
         if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-            # 判断是否是原图 (文件名包含 _origin)
             if "_origin" in filename.lower():
                 groups[prefix]["origin"] = file_path
             else:
                 groups[prefix]["candidates"].append(file_path)
-                
-        # 文本处理
         elif filename.lower().endswith('.txt'):
             try:
                 with open(file_path, "r", encoding="utf-8") as f:
@@ -57,7 +52,6 @@ def load_data():
                 with open(file_path, "r", encoding="gbk") as f:
                     groups[prefix]["instruction"] = f.read()
 
-    # 过滤掉既没有原图也没有候选图的组
     valid_groups = {}
     for k, v in groups.items():
         if v["origin"] is not None or len(v["candidates"]) > 0:
@@ -70,18 +64,11 @@ 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
-    
-    row = [
-        user_id,
-        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-        group_id,
-        choice_labels, 
-        method_names    
-    ]
+    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()
@@ -90,216 +77,186 @@ def save_user_vote(user_id, group_id, choice_labels, method_names):
             if not file_exists:
                 writer.writerow(["user_id", "timestamp", "group_id", "selected_labels", "selected_methods"])
             writer.writerow(row)
-    
-    print(f"Saved: User {user_id} selected {method_names}")
+    print(f"Saved: {user_id} -> {choice_labels}")
 
-# --- 5. 交互逻辑 ---
+# --- 5. 核心逻辑：状态管理 ---
 
-def get_current_question_ui(user_state):
-    """刷新界面：原图置顶 + 候选图乱序"""
-    current_idx = user_state["index"]
-    
-    if current_idx >= len(ALL_GROUP_IDS):
-        return (
-            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), 
-            gr.update(visible=False), gr.update(value="## 🎉 测试结束！\n感谢您的参与。", visible=True),
-            user_state, [], [], False # False表示当前没有原图显示
-        )
+def init_state():
+    """初始化用户状态"""
+    return {
+        "user_id": str(uuid.uuid4())[:8],
+        "index": 0,
+        "is_finished": False
+    }
+
+def next_question_data(user_state):
+    """
+    计算下一题的数据，并返回给前端渲染。
+    这里只处理数据准备（打乱顺序等），不处理 UI。
+    """
+    idx = user_state["index"]
     
-    group_id = ALL_GROUP_IDS[current_idx]
+    if idx >= len(ALL_GROUP_IDS):
+        user_state["is_finished"] = True
+        return user_state, None, [], [] # 结束状态
+
+    group_id = ALL_GROUP_IDS[idx]
     group_data = ALL_GROUPS[group_id]
     
-    instruction_text = f"### 任务 ({current_idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
+    # 准备原图
+    origin_path = group_data["origin"]
     
-    # --- 构建显示列表 ---
-    display_list = []
-    real_paths = [] # 记录真实的路径顺序
-    
-    # 1. 先放原图 (如果有)
-    has_origin = False
-    if group_data["origin"]:
-        display_list.append((group_data["origin"], "Reference (原图)"))
-        real_paths.append(group_data["origin"])
-        has_origin = True
-        
-    # 2. 再放候选图 (打乱顺序)
+    # 准备候选图（打乱）
     candidates = group_data["candidates"].copy()
     random.shuffle(candidates)
     
-    for i, img_path in enumerate(candidates):
-        # 标签从 Option A 开始
-        label = f"Option {chr(65+i)}" 
-        display_list.append((img_path, label))
-        real_paths.append(img_path)
-    
-    # 动态计算列数 (原图+候选图的总数)
-    total_imgs = len(display_list)
-    # 如果总数是2(1原图+1候选)，显示2列
-    # 如果总数是3，显示3列
-    # 如果总数是4或以上，显示3列换行，或者你可以改成 4
-    cols = min(total_imgs, 3) 
+    # 构造候选图信息列表 [(path, "Option A"), (path, "Option B")...]
+    candidate_info = []
+    for i, path in enumerate(candidates):
+        label = f"Option {chr(65+i)}"
+        candidate_info.append({"path": path, "label": label})
 
-    return (
-        gr.update(value=instruction_text, visible=True),
-        gr.update(value=display_list, columns=cols, visible=True), 
-        gr.update(value="当前未选择任何图片", visible=True), 
-        gr.update(visible=True), 
-        gr.update(visible=False), 
-        user_state,
-        real_paths, # 存下当前界面所有图片的真实路径
-        [],         # 清空选中项
-        has_origin  # 告诉前端第一张是不是原图
-    )
+    return user_state, origin_path, candidate_info, group_data["instruction"]
 
-def toggle_selection(evt: gr.SelectData, current_indices, has_origin):
-    """
-    处理点击。
-    has_origin: 如果为 True，说明 Index 0 是原图，禁止选中。
-    """
-    clicked_idx = evt.index
-    
-    # --- 保护机制：如果第一张是原图，点击无效 ---
-    if has_origin and clicked_idx == 0:
-        # 你可以去掉这个if，允许选原图，但通常 User Study 不需要选 Reference
-        # 这里我们返回原样，不更新选中状态
-        return current_indices, "⚠️ 您点击的是参考原图 (Reference)，请选择后面的选项。"
-    
-    # 切换选中状态
-    if clicked_idx in current_indices:
-        current_indices.remove(clicked_idx)
-    else:
-        current_indices.append(clicked_idx)
-    current_indices.sort()
-    
-    # 生成状态文本
-    if not current_indices:
-        status_text = "当前未选择任何图片"
-    else:
-        # 计算显示的 Label
-        labels = []
-        for idx in current_indices:
-            # 如果有原图，index 1 才是 Option A
-            if has_origin:
-                option_char = chr(65 + (idx - 1)) # idx=1 -> A, idx=2 -> B
-            else:
-                option_char = chr(65 + idx)
-            labels.append(f"Option {option_char}")
-            
-        status_text = "已选中: " + ", ".join(labels)
-    
-    return current_indices, status_text
+def submit_logic(user_state, current_candidates, selected_indices, is_none=False):
+    """处理提交，保存数据，并进入下一题"""
+    if user_state["is_finished"]:
+        return user_state, [], []
 
-def submit_vote(user_state, current_file_paths, current_indices, is_none=False, has_origin=False):
-    user_id = user_state["user_id"]
     current_idx = user_state["index"]
-    
-    if current_idx >= len(ALL_GROUP_IDS):
-        return get_current_question_ui(user_state)
-    
     group_id = ALL_GROUP_IDS[current_idx]
-    
+
+    # 1. 保存数据
     if is_none:
-        save_user_vote(user_id, group_id, "Rejected All", "None_Satisfied")
-        user_state["index"] += 1
-        return get_current_question_ui(user_state)
-    
-    if not current_indices:
-        return (
-            gr.update(), gr.update(), 
-            gr.update(value="❌ 请至少选择一张图片"), 
-            gr.update(), gr.update(), 
-            user_state, current_file_paths, current_indices, has_origin
-        )
-    
-    selected_labels = []
-    selected_methods = []
-    
-    for idx in current_indices:
-        # 获取真实路径
-        real_path = current_file_paths[idx]
-        
-        # 计算 Label (Option A/B...)
-        if has_origin:
-            label = f"Option {chr(65 + (idx - 1))}"
-        else:
-            label = f"Option {chr(65 + idx)}"
-        selected_labels.append(label)
-        
-        # 提取方法名
-        filename = os.path.basename(real_path)
-        name_no_ext = os.path.splitext(filename)[0]
-        
-        # 如果是原图被选中了(假设你去掉了保护机制)，方法名就是 origin
-        if "_origin" in name_no_ext:
-            method = "reference_origin"
-        else:
+        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)
             method = parts[1] if len(parts) > 1 else name_no_ext
+            methods.append(method)
             
-        selected_methods.append(method)
-    
-    str_labels = "; ".join(selected_labels)
-    str_methods = "; ".join(selected_methods)
-    
-    save_user_vote(user_id, group_id, str_labels, str_methods)
-    
+        save_user_vote(user_state["user_id"], group_id, "; ".join(labels), "; ".join(methods))
+
+    # 2. 索引+1
     user_state["index"] += 1
-    return get_current_question_ui(user_state)
+    
+    # 3. 清空选择
+    return user_state, [] # 返回新的 state 和空的 selected_indices
 
-# --- 6. 界面构建 ---
+# --- 6. 界面构建 (使用 @gr.render 实现自定义布局) ---
 with gr.Blocks(title="User Study") as demo:
     
-    state_user = gr.State(lambda: {"user_id": str(uuid.uuid4())[:8], "index": 0})
-    state_files = gr.State([])      
-    state_indices = gr.State([]) 
-    state_has_origin = gr.State(False) # 记录当前页面是否有原图
+    # === 状态变量 ===
+    # state_main: 存 user_id, index
+    state_main = gr.State(init_state())
+    # state_current_data: 存当前题目的数据 (原图, 候选图列表, 说明文)
+    state_origin = gr.State()
+    state_candidates = gr.State([]) # list of dict
+    state_instruction = gr.State("")
+    # state_selection: 存当前选中的索引 [0, 2]
+    state_selection = gr.State([])
+
+    # === 页面布局 ===
+    header = gr.Markdown("Loading...")
     
-    with gr.Column():
-        instruction_md = gr.Markdown("Loading...")
+    # 动态渲染区域
+    @gr.render(inputs=[state_main, state_origin, state_candidates, state_selection, state_instruction])
+    def render_content(main_st, origin, candidates, selection, instruction):
         
-        # 图片显示
-        gallery = gr.Gallery(
-            label="请参考第一张原图，选择后面最好的结果（可多选）", 
-            allow_preview=True, 
-            object_fit="contain", 
-            height="auto",
-            interactive=True
-        )
-        
-        status_box = gr.Textbox(value="当前未选择任何图片", label="当前选中状态", interactive=False)
-        
-        with gr.Row():
-            btn_submit = gr.Button("✅ 提交选择", variant="primary", scale=2)
-            btn_none = gr.Button("🚫 都不满意", variant="stop", scale=1)
+        # 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}")
         
-        end_msg = gr.Markdown(visible=False)
+        # 3. 显示原图 (Reference)
+        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**点击图片下方的按钮进行选择。**")
 
-    # 加载第一题
-    demo.load(
-        fn=get_current_question_ui, 
-        inputs=[state_user], 
-        outputs=[instruction_md, gallery, status_box, btn_submit, end_msg, state_user, state_files, state_indices, state_has_origin]
-    )
-    
-    # 点击图片
-    gallery.select(
-        fn=toggle_selection,
-        inputs=[state_indices, state_has_origin],
-        outputs=[state_indices, status_box]
-    )
+        # 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 效果
+                    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
+                    
+                    # 绑定点击事件，更新 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=lambda s, f, i, h: submit_vote(s, f, i, is_none=False, has_origin=h),
-        inputs=[state_user, state_files, state_indices, state_has_origin],
-        outputs=[instruction_md, gallery, status_box, btn_submit, end_msg, state_user, state_files, state_indices, state_has_origin]
+        fn=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=lambda s, f, i, h: submit_vote(s, f, i, is_none=True, has_origin=h),
-        inputs=[state_user, state_files, state_indices, state_has_origin],
-        outputs=[instruction_md, gallery, status_box, btn_submit, end_msg, state_user, state_files, state_indices, state_has_origin]
+        fn=on_none,
+        inputs=[state_main, state_candidates, state_selection],
+        outputs=[state_main, state_selection, state_origin, state_candidates, state_instruction]
     )
 
 if __name__ == "__main__":