Update app.py

app.py

@@ -8,213 +8,276 @@ from pathlib import Path
 from huggingface_hub import CommitScheduler
 
 # --- 1. 配置区域 ---
-
-# 你的数据集地址 (已修改)
-DATASET_REPO_ID = "Emilyxml/moveit"
-
-# 数据源文件夹 (读取你的图片和txt)
-DATA_FOLDER = "data"
-
-# 临时日志文件夹 (用于存放用户生成的CSV，Scheduler 会监控这里)
-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 (需要在 Space 设置里配置 HF_TOKEN)
 TOKEN = os.environ.get("HF_TOKEN")
 
 # --- 2. 启动同步调度器 ---
-# 只要 logs 文件夹里有 CSV 变化，就自动上传到 Dataset 的 data 文件夹下
 scheduler = CommitScheduler(
     repo_id=DATASET_REPO_ID,
     repo_type="dataset",
-    folder_path=LOG_FOLDER, # 监控本地的 logs 文件夹
-    path_in_repo="data",    # 上传到 Dataset 的 data 目录中
-    every=1,                # 每分钟同步一次
+    folder_path=LOG_FOLDER,
+    path_in_repo="data",
+    every=1,
     token=TOKEN
 )
 
-# --- 3. 数据加载逻辑 ---
+# --- 3. 数据加载逻辑 (保持不变) ---
 def load_data():
     groups = {}
-    
-    # 检查 data 文件夹是否存在
     if not os.path.exists(DATA_FOLDER):
-        # 如果不存在，尝试创建（防止报错），但实际应该由你上传文件
         os.makedirs(DATA_FOLDER, exist_ok=True)
-        print(f"Warning: {DATA_FOLDER} not found. Please upload your images.")
         return {}, []
 
-    # 遍历文件
     for filename in os.listdir(DATA_FOLDER):
-        if filename.startswith('.'): continue # 跳过隐藏文件
-        
+        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:
             groups[prefix] = {"images": [], "instruction": "暂无说明"}
 
-        # 识别图片
-        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp')):
+        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
             groups[prefix]["images"].append(file_path)
-        # 识别文本
         elif filename.lower().endswith('.txt'):
             try:
                 with open(file_path, "r", encoding="utf-8") as f:
                     groups[prefix]["instruction"] = f.read()
             except:
-                # 兼容 gbk 编码
                 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}
-    
-    # 生成题目列表，并随机打乱
     group_ids = list(valid_groups.keys())
     random.shuffle(group_ids)
-    
-    print(f"Loaded {len(group_ids)} groups of images.")
+    print(f"Loaded {len(group_ids)} groups.")
     return valid_groups, group_ids
 
-# 全局加载数据
 ALL_GROUPS, ALL_GROUP_IDS = load_data()
 
-# --- 4. 保存逻辑 (每个用户一个独立CSV) ---
-def save_user_vote(user_id, group_id, choice_label, method_name):
+# --- 4. 保存逻辑 (支持多选保存) ---
+def save_user_vote(user_id, group_id, choice_labels, method_names):
     """
-    保存单次投票到 logs/user_{user_id}.csv
+    保存投票。
+    choice_labels: 字符串，例如 "Option A; Option B"
+    method_names: 字符串，例如 "omnigen; sdxl"
     """
     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_label,  # 用户选了 Option A 还是 B
-        method_name    # 真实的方法名
+        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_label", "selected_method"])
+                writer.writerow(["user_id", "timestamp", "group_id", "selected_labels", "selected_methods"])
             writer.writerow(row)
     
-    print(f"Saved vote for {user_id}: {method_name}")
+    print(f"Saved: User {user_id} selected {method_names}")
+
+# --- 5. 交互逻辑 (多选核心) ---
 
-# --- 5. 交互逻辑 ---
-def get_next_question(user_state):
+def get_current_question_ui(user_state):
+    """根据当前索引刷新界面"""
     current_idx = user_state["index"]
     
-    # 1. 检查是否做完
+    # 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(visible=False),
-            gr.update(visible=False),
-            gr.update(value="## 🎉 测试结束！\n感谢您的参与，您的选择已保存。", visible=True),
+            gr.update(value="## 🎉 测试结束！\n感谢您的参与，所有结果已保存。", visible=True),
             user_state,
-            []
+            [], # 清空文件路径
+            []  # 清空当前选中的索引
         )
     
-    # 2. 获取当前组
+    # 2. 获取数据
     group_id = ALL_GROUP_IDS[current_idx]
     group_data = ALL_GROUPS[group_id]
     
-    # 3. 准备 Prompt
-    instruction_text = f"## 任务 ({current_idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
+    # 3. 准备文本
+    instruction_text = f"### 任务 ({current_idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
     
-    # 4. 准备图片 (打乱顺序实现盲测)
+    # 4. 准备图片 (盲测 + 打乱)
     original_images = group_data["images"]
     shuffled_images = original_images.copy()
     random.shuffle(shuffled_images)
     
-    # 构造 Gradio 显示对象
+    # 构造显示列表
     display_list = []
     for i, img_path in enumerate(shuffled_images):
         label = f"Option {chr(65+i)}" # Option A, Option B...
         display_list.append((img_path, label))
-        
+    
+    # 动态列数
+    num_imgs = len(shuffled_images)
+    cols = 2 if num_imgs == 4 else min(num_imgs, 3)
+
     return (
         gr.update(value=instruction_text, visible=True),
-        gr.update(value=display_list, visible=True),
-        gr.update(visible=True),
-        gr.update(visible=False),
+        gr.update(value=display_list, columns=cols, visible=True), # Gallery
+        gr.update(value="当前未选择任何图片", visible=True), # 状态栏重置
+        gr.update(visible=True), # 按钮区可见
+        gr.update(visible=False), # 结束语
         user_state,
-        shuffled_images # 将乱序后的真实路径列表传给 State
+        shuffled_images,
+        [] # 重置选中的索引列表
     )
 
-def on_vote(user_state, current_file_paths, select_data: gr.SelectData = None, is_none=False):
+def toggle_selection(evt: gr.SelectData, current_indices):
+    """
+    处理图片点击：
+    点击一次 -> 选中
+    再点一次 -> 取消选中
+    """
+    clicked_idx = evt.index
+    
+    # 切换状态
+    if clicked_idx in current_indices:
+        current_indices.remove(clicked_idx)
+    else:
+        current_indices.append(clicked_idx)
+        
+    # 排序一下，让显示更好看 (Option A, Option B)
+    current_indices.sort()
+    
+    # 更新状态文本
+    if not current_indices:
+        status_text = "当前未选择任何图片"
+    else:
+        labels = [f"Option {chr(65+i)}" for i in current_indices]
+        status_text = "已选中: " + ", ".join(labels)
+    
+    return current_indices, status_text
+
+def submit_vote(user_state, current_file_paths, current_indices, is_none=False):
+    """
+    提交投票（可能是多选，可能是None）
+    """
     user_id = user_state["user_id"]
     current_idx = user_state["index"]
     
-    # 防止溢出
     if current_idx >= len(ALL_GROUP_IDS):
-        return get_next_question(user_state)
-        
+        return get_current_question_ui(user_state)
+    
     group_id = ALL_GROUP_IDS[current_idx]
-    selected_method = "Unknown"
-    selected_label = "None"
     
-    # --- 解析选择 ---
+    # --- 场景1: 都没有 ---
     if is_none:
-        selected_method = "None_Satisfied"
-        selected_label = "Rejected All"
-    elif select_data is not None:
-        idx = select_data.index
-        # 获取真实路径
-        real_image_path = current_file_paths[idx]
-        selected_label = select_data.value["caption"]
+        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(), gr.update(), 
+            user_state, current_file_paths, current_indices
+        )
+    
+    # 解析所有选中的图片
+    selected_labels = []
+    selected_methods = []
+    
+    for idx in current_indices:
+        # 1. 记录 Option X
+        label = f"Option {chr(65+idx)}"
+        selected_labels.append(label)
         
-        # 提取方法名 (例如 6180a_omnigen.png -> omnigen)
-        filename = os.path.basename(real_image_path)
+        # 2. 提取方法名
+        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)
-        if len(parts) > 1:
-            selected_method = parts[1]
-        else:
-            selected_method = name_no_ext
-            
-    # --- 保存 ---
-    save_user_vote(user_id, group_id, selected_label, selected_method)
-    
-    # --- 下一题 ---
+        method = parts[1] if len(parts) > 1 else name_no_ext
+        selected_methods.append(method)
+    
+    # 用分号连接 (CSV友好)
+    str_labels = "; ".join(selected_labels)
+    str_methods = "; ".join(selected_methods)
+    
+    save_user_vote(user_id, group_id, str_labels, str_methods)
+    
+    # 下一题
     user_state["index"] += 1
-    return get_next_question(user_state)
+    return get_current_question_ui(user_state)
 
 # --- 6. 界面构建 ---
-with gr.Blocks(title="User Study - MoveIt") as demo:
-    # State 初始化：每次刷新网页生成新的 user_id
+with gr.Blocks(title="Multi-Select User Study", theme=gr.themes.Soft()) as demo:
+    
+    # 状态变量
     state_user = gr.State(lambda: {"user_id": str(uuid.uuid4())[:8], "index": 0})
-    state_files = gr.State([])
+    state_files = gr.State([])      # 存当前图片的真实路径
+    state_indices = gr.State([])    # 存当前选中的图片索引 [0, 2]
     
-    with gr.Column(elem_id="main"):
+    with gr.Column():
         instruction_md = gr.Markdown("Loading...")
         
+        # 图片区
         gallery = gr.Gallery(
-            label="请点击选择最佳图片",
-            columns=[2],
+            label="请点击选择图片（可多选）", 
+            allow_preview=True, 
+            object_fit="contain", 
             height="auto",
             interactive=True
         )
         
-        btn_none = gr.Button("🚫 没有任何一张图片符合要求", variant="stop")
+        # 状态显示区（告诉用户选了啥）
+        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)
         
         end_msg = gr.Markdown(visible=False)
 
-    # 事件绑定
-    demo.load(fn=get_next_question, inputs=[state_user], outputs=[instruction_md, gallery, btn_none, end_msg, state_user, state_files])
+    # --- 事件流 ---
+    
+    # 1. 启动加载
+    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]
+    )
     
-    gallery.select(fn=lambda s, f, evt: on_vote(s, f, evt, is_none=False), inputs=[state_user, state_files], outputs=[instruction_md, gallery, btn_none, end_msg, state_user, state_files])
+    # 2. 点击图片 -> 切换选中状态 (不翻页)
+    gallery.select(
+        fn=toggle_selection,
+        inputs=[state_indices],
+        outputs=[state_indices, status_box]
+    )
     
-    btn_none.click(fn=lambda s, f: on_vote(s, f, None, is_none=True), inputs=[state_user, state_files], outputs=[instruction_md, gallery, btn_none, end_msg, state_user, state_files])
+    # 3. 点击提交 -> 保存并下一页
+    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]
+    )
+    
+    # 4. 点击都不满意 -> 保存并下一页
+    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]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()