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] 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}) 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) 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) # --- 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] ) if __name__ == "__main__": demo.launch()