Update app.py

app.py

@@ -1,252 +1,52 @@
-import gradio as gr
-import os
-import random
-import uuid
-import csv
-from datetime import datetime
-from pathlib import Path
-from PIL import Image
-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"]
+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 格式
         )
-        print("✅ 数据下载完成！")
-    except Exception as e:
-        print(f"⚠️ 下载失败: {e}")
-
-# --- 3. 恢复后台同步 (解决点击卡顿的关键) ---
-# 使用 Scheduler，提交操作不需要等待网络上传，速度最快
-scheduler = CommitScheduler(
-    repo_id=DATASET_REPO_ID,
-    repo_type="dataset",
-    folder_path=LOG_FOLDER,
-    path_in_repo="logs",
-    every=1, # 每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]
         
-        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=500):
-    """
-    调整大小至 500px，对于 User Study 的缩略图查看完全足够。
-    """
-    if not image_path:
-        return None
-    try:
-        img = Image.open(image_path)
-        # 转换为 RGB 防止 PNG 透明通道在 JPEG 转换时报错
-        if img.mode in ("RGBA", "P"):
-            img = img.convert("RGB")
-            
-        if img.width > max_width:
-            ratio = max_width / img.width
-            new_height = int(img.height * ratio)
-            # 使用 LANCZOS 算法保证缩放质量
-            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, []
+        gr.Markdown("👇 **请在下方勾选您认为最好的结果（可多选）：**")
+        
+        checkbox_options = gr.CheckboxGroup(
+            choices=[], 
+            label="您的选择", 
+            info="对应上方图片的标签 (Option A, B...)"
         )
-    
-    group_id = ALL_GROUP_IDS[idx]
-    group_data = ALL_GROUPS[group_id]
-    
-    # 1. 优化原图
-    origin_img = optimize_image(group_data["origin"], max_width=500)
-    
-    # 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=500)
-        gallery_items.append((optimized_img, 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_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
-    )
-
-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)
+        with gr.Row():
+            btn_submit = gr.Button("🚀 提交 (Submit)", variant="primary")
+            btn_none = gr.Button("🚫 都不满意 (None)", variant="stop")
 
-    # --- 极速保存：只写本地文件，不等待网络上传 ---
-    user_filename = f"user_{user_state['user_id']}.csv"
-    user_file_path = LOG_FOLDER / user_filename
-    
-    # 使用 Scheduler 提供的锁来保证多线程安全
-    with scheduler.lock:
-        file_exists = user_file_path.exists()
-        with open(user_file_path, "a", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            if not file_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
-            ])
-    
-    print(f"✅ Local Saved: {group_id} (Upload will happen in background)")
-            
-    user_state["index"] += 1
-    return get_next_question(user_state)
+md_end = gr.Markdown(visible=False)
 
-# --- 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):
-            # 强制 JPEG 格式，quality 默认 90，显示速度快
-            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"
-            )
-            gr.Markdown("👇 **请在下方勾选您认为最好的结果（可多选）：**")
-            checkbox_options = gr.CheckboxGroup(choices=[], label="您的选择", info="对应上方图片的标签")
-            
-            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]
+)
 
-    # 初始加载
-    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_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]
+)
 
-if __name__ == "__main__":
-    demo.launch()
+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]
+)