Update app.py

app.py

@@ -24,7 +24,7 @@ scheduler = CommitScheduler(
     token=TOKEN
 )
 
-# --- 3. 数据加载逻辑 ---
+# --- 3. 数据加载逻辑 (区分原图和候选图) ---
 def load_data():
     groups = {}
     if not os.path.exists(DATA_FOLDER):
@@ -34,13 +34,21 @@ 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]
+        prefix = filename[:5] # 以前5个字符作为组ID
         
         if prefix not in groups:
-            groups[prefix] = {"images": [], "instruction": "暂无说明"}
+            # origin 存原图路径，candidates 存其他方法的图
+            groups[prefix] = {"origin": None, "candidates": [], "instruction": "暂无说明"}
 
+        # 图片处理
         if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-            groups[prefix]["images"].append(file_path)
+            # 判断是否是原图 (文件名包含 _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:
@@ -49,7 +57,12 @@ def load_data():
                 with open(file_path, "r", encoding="gbk") as f:
                     groups[prefix]["instruction"] = f.read()
 
-    valid_groups = {k: v for k, v in groups.items() if len(v["images"]) > 0}
+    # 过滤掉既没有原图也没有候选图的组
+    valid_groups = {}
+    for k, v in groups.items():
+        if v["origin"] is not None or len(v["candidates"]) > 0:
+            valid_groups[k] = v
+            
     group_ids = list(valid_groups.keys())
     random.shuffle(group_ids)
     print(f"Loaded {len(group_ids)} groups.")
@@ -83,43 +96,48 @@ def save_user_vote(user_id, group_id, choice_labels, method_names):
 # --- 5. 交互逻辑 ---
 
 def get_current_question_ui(user_state):
-    """根据当前索引刷新界面"""
+    """刷新界面：原图置顶 + 候选图乱序"""
     current_idx = user_state["index"]
     
-    # 1. 检查是否结束
     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,
-            [], 
-            []  
+            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表示当前没有原图显示
         )
     
-    # 2. 获取数据
     group_id = ALL_GROUP_IDS[current_idx]
     group_data = ALL_GROUPS[group_id]
     
-    # 3. 准备文本
     instruction_text = f"### 任务 ({current_idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
     
-    # 4. 准备图片
-    original_images = group_data["images"]
-    shuffled_images = original_images.copy()
-    random.shuffle(shuffled_images)
-    
-    # 构造显示列表
+    # --- 构建显示列表 ---
     display_list = []
-    for i, img_path in enumerate(shuffled_images):
+    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)
     
-    # 动态列数
-    num_imgs = len(shuffled_images)
-    cols = 2 if num_imgs == 4 else min(num_imgs, 3)
+    # 动态计算列数 (原图+候选图的总数)
+    total_imgs = len(display_list)
+    # 如果总数是2(1原图+1候选)，显示2列
+    # 如果总数是3，显示3列
+    # 如果总数是4或以上，显示3列换行，或者你可以改成 4
+    cols = min(total_imgs, 3) 
 
     return (
         gr.update(value=instruction_text, visible=True),
@@ -128,29 +146,50 @@ def get_current_question_ui(user_state):
         gr.update(visible=True), 
         gr.update(visible=False), 
         user_state,
-        shuffled_images,
-        [] 
+        real_paths, # 存下当前界面所有图片的真实路径
+        [],         # 清空选中项
+        has_origin  # 告诉前端第一张是不是原图
     )
 
-def toggle_selection(evt: gr.SelectData, current_indices):
+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:
-        labels = [f"Option {chr(65+i)}" for i in current_indices]
+        # 计算显示的 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_vote(user_state, current_file_paths, current_indices, is_none=False):
+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"]
     
@@ -159,33 +198,44 @@ def submit_vote(user_state, current_file_paths, current_indices, is_none=False):
     
     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)
     
-    # 场景2: 提交多选
     if not current_indices:
         return (
             gr.update(), gr.update(), 
-            gr.update(value="❌ 请至少选择一张图片，或者点击“都不满意”"), 
+            gr.update(value="❌ 请至少选择一张图片"), 
             gr.update(), gr.update(), 
-            user_state, current_file_paths, current_indices
+            user_state, current_file_paths, current_indices, has_origin
         )
     
     selected_labels = []
     selected_methods = []
     
     for idx in current_indices:
-        label = f"Option {chr(65+idx)}"
+        # 获取真实路径
+        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)
         
-        real_path = current_file_paths[idx]
+        # 提取方法名
         filename = os.path.basename(real_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
+        
+        # 如果是原图被选中了(假设你去掉了保护机制)，方法名就是 origin
+        if "_origin" in name_no_ext:
+            method = "reference_origin"
+        else:
+            parts = name_no_ext.split('_', 1)
+            method = parts[1] if len(parts) > 1 else name_no_ext
+            
         selected_methods.append(method)
     
     str_labels = "; ".join(selected_labels)
@@ -196,18 +246,20 @@ def submit_vote(user_state, current_file_paths, current_indices, is_none=False):
     user_state["index"] += 1
     return get_current_question_ui(user_state)
 
-# --- 6. 界面构建 (已移除 theme 参数) ---
-with gr.Blocks(title="Multi-Select User Study") as demo:
+# --- 6. 界面构建 ---
+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_indices = gr.State([]) 
+    state_has_origin = gr.State(False) # 记录当前页面是否有原图
     
     with gr.Column():
         instruction_md = gr.Markdown("Loading...")
         
+        # 图片显示
         gallery = gr.Gallery(
-            label="请点击选择图片（可多选）", 
+            label="请参考第一张原图，选择后面最好的结果（可多选）", 
             allow_preview=True, 
             object_fit="contain", 
             height="auto",
@@ -217,33 +269,37 @@ with gr.Blocks(title="Multi-Select User Study") as demo:
         status_box = gr.Textbox(value="当前未选择任何图片", label="当前选中状态", interactive=False)
         
         with gr.Row():
-            btn_submit = gr.Button("✅ 提交选择 (Confirm Selection)", variant="primary", scale=2)
-            btn_none = gr.Button("🚫 都不满意 (None of them)", variant="stop", scale=1)
+            btn_submit = gr.Button("✅ 提交选择", variant="primary", scale=2)
+            btn_none = gr.Button("🚫 都不满意", variant="stop", scale=1)
         
         end_msg = gr.Markdown(visible=False)
 
+    # 加载第一题
     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]
+        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],
+        inputs=[state_indices, state_has_origin],
         outputs=[state_indices, status_box]
     )
     
+    # 提交按钮
     btn_submit.click(
-        fn=lambda s, f, i: submit_vote(s, f, i, is_none=False),
-        inputs=[state_user, state_files, state_indices],
-        outputs=[instruction_md, gallery, status_box, btn_submit, end_msg, state_user, state_files, state_indices]
+        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]
     )
     
+    # 都不满意按钮
     btn_none.click(
-        fn=lambda s, f, i: submit_vote(s, f, i, is_none=True),
-        inputs=[state_user, state_files, state_indices],
-        outputs=[instruction_md, gallery, status_box, btn_submit, end_msg, state_user, state_files, state_indices]
+        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]
     )
 
 if __name__ == "__main__":