Update app.py

app.py

@@ -1,52 +1,259 @@
-state_user = gr.State(lambda: {"user_id": str(uuid.uuid4())[:8], "index": 0})
-state_candidates_info = gr.State([]) 
-
-with gr.Row():
-    md_instruction = gr.Markdown("Loading...")
-
-with gr.Row():
-    with gr.Column(scale=1):
-        # 将 format 设置为 jpeg 进一步减小体积
-        img_origin = gr.Image(label="Reference (参考原图)", interactive=False, height=400, format="jpeg")
-    
-    with gr.Column(scale=2):
-        gallery_candidates = gr.Gallery(
-            label="Candidates (候选结果)", 
-            columns=[2], 
-            height="auto",
-            object_fit="contain",
-            interactive=False,
-            format="jpeg" # 强制输出 JPEG 格式
+import gradio as gr
+import os
+import random
+import uuid
+import csv
+from datetime import datetime
+from pathlib import Path
+from PIL import Image  # 引入 PIL 用于处理图片
+from huggingface_hub import CommitScheduler, snapshot_download
+
+# --- 1. 配置区域 ---
+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. 自动下载数据 ---
+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"⚠️ 下载失败: {e}")
+
+# --- 3. 启动同步调度器 ---
+scheduler = CommitScheduler(
+    repo_id=DATASET_REPO_ID,
+    repo_type="dataset",
+    folder_path=LOG_FOLDER,
+    path_in_repo="logs",
+    every=1,
+    token=TOKEN
+)
+
+# --- 4. 数据加载 ---
+def load_data():
+    groups = {}
+    if not os.path.exists(DATA_FOLDER):
+        return {}, []
+
+    for filename in os.listdir(DATA_FOLDER):
+        if filename.startswith('.'): continue
+        file_path = os.path.join(DATA_FOLDER, filename)
+        prefix = filename[:5]
         
-        gr.Markdown("👇 **请在下方勾选您认为最好的结果（可多选）：**")
-        
-        checkbox_options = gr.CheckboxGroup(
-            choices=[], 
-            label="您的选择", 
-            info="对应上方图片的标签 (Option A, B...)"
+        if prefix not in groups:
+            groups[prefix] = {"origin": None, "candidates": [], "instruction": "暂无说明"}
+
+        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+            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:
+                    groups[prefix]["instruction"] = f.read()
+            except:
+                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:
+            valid_groups[k] = v
+            
+    group_ids = list(valid_groups.keys())
+    random.shuffle(group_ids)
+    print(f"Loaded {len(group_ids)} groups.")
+    return valid_groups, group_ids
+
+ALL_GROUPS, ALL_GROUP_IDS = load_data()
+
+# --- NEW: 图片优化函数 (提速关键) ---
+def optimize_image(image_path, max_width=800):
+    """
+    读取图片并调整大小，减少传输时间。
+    max_width: 限制最大宽度为 800px (足够人眼评估)
+    """
+    if not image_path:
+        return None
+    try:
+        img = Image.open(image_path)
+        # 如果图片太大，就缩小
+        if img.width > max_width:
+            ratio = max_width / img.width
+            new_height = int(img.height * ratio)
+            img = img.resize((max_width, new_height), Image.LANCZOS)
+        return img
+    except Exception as e:
+        print(f"Error loading image {image_path}: {e}")
+        return None
+
+# --- 5. 核心逻辑 ---
+
+def get_next_question(user_state):
+    """准备下一题的数据"""
+    idx = user_state["index"]
+    
+    if 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="## 🎉 测试结束！感谢您的参与。", visible=True),
+            user_state,
+            [] 
         )
+    
+    group_id = ALL_GROUP_IDS[idx]
+    group_data = ALL_GROUPS[group_id]
+    
+    # 1. 优化原图 (返回 PIL 对象而不是路径)
+    origin_img = optimize_image(group_data["origin"], max_width=600)
+    
+    # 2. 优化候选图
+    candidates = group_data["candidates"].copy()
+    random.shuffle(candidates)
+    
+    gallery_items = []
+    choices = []
+    candidates_info = [] 
+    
+    for i, path in enumerate(candidates):
+        label = f"Option {chr(65+i)}"
+        
+        # 优化每张候选图
+        optimized_img = optimize_image(path, max_width=600)
+        
+        gallery_items.append((optimized_img, label))
+        choices.append(label)
+        candidates_info.append({"label": label, "path": path})
         
-        with gr.Row():
-            btn_submit = gr.Button("🚀 提交 (Submit)", variant="primary")
-            btn_none = gr.Button("🚫 都不满意 (None)", variant="stop")
+    instruction = f"### 任务 ({idx + 1} / {len(ALL_GROUP_IDS)})\n\n{group_data['instruction']}"
+    
+    return (
+        gr.update(value=origin_img, visible=True if origin_img 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
+    )
 
-md_end = gr.Markdown(visible=False)
+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:
+        choice_str = "Rejected All"
+        method_str = "None_Satisfied"
+    else:
+        if not selected_options:
+            raise gr.Error("请至少勾选一个选项，或点击“都不满意”")
+            
+        choice_str = "; ".join(selected_options)
+        selected_methods = []
+        for opt in selected_options: 
+            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)
 
-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]
-)
+    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 get_next_question(user_state)
 
-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]
-)
+# --- 6. 界面构建 ---
+with gr.Blocks(title="User Study") as demo:
+    
+    state_user = gr.State(lambda: {"user_id": str(uuid.uuid4())[:8], "index": 0})
+    state_candidates_info = gr.State([]) 
+    
+    with gr.Row():
+        md_instruction = gr.Markdown("Loading...")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # 将 format 设置为 jpeg 进一步减小体积
+            img_origin = gr.Image(label="Reference (参考原图)", interactive=False, height=400, format="jpeg")
+        
+        with gr.Column(scale=2):
+            gallery_candidates = gr.Gallery(
+                label="Candidates (候选结果)", 
+                columns=[2], 
+                height="auto",
+                object_fit="contain",
+                interactive=False,
+                format="jpeg" # 强制输出 JPEG 格式
+            )
+            
+            gr.Markdown("👇 **请在下方勾选您认为最好的结果（可多选）：**")
+            
+            checkbox_options = gr.CheckboxGroup(
+                choices=[], 
+                label="您的选择", 
+                info="对应上方图片的标签 (Option A, B...)"
+            )
+            
+            with gr.Row():
+                btn_submit = gr.Button("🚀 提交 (Submit)", variant="primary")
+                btn_none = gr.Button("🚫 都不满意 (None)", variant="stop")
+    
+    md_end = gr.Markdown(visible=False)
+
+    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]
+    )
+    
+    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]
+    )
+    
+    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]
+    )
 
-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()