Update app.py

app.py

@@ -2,14 +2,53 @@ import os
 import json
 import gradio as gr
 from pathlib import Path
+from io import BytesIO
+import base64
+import requests
+from PIL import Image
 
 # =============== 配置：HF dataset 仓库 ID =================
 HF_DATASET_ID = "maplebb/UniREditBench-Results"
 HF_BASE_URL = f"https://huggingface.co/datasets/{HF_DATASET_ID}/resolve/main"
+BASE_HF_URL = f"https://huggingface.co/datasets/{HF_DATASET_ID}/resolve/main"
+
+_image_cache = {}
+_html_cache = {}
+
+def get_url_response(url):
+    try:
+        resp = requests.get(url, timeout=10)
+        resp.raise_for_status()
+        return resp
+    except Exception as e:
+        print(f"Error fetching {url}: {e}")
+        return None
+
+def load_image_uniredit(rel_path: str):
+    """从 UniREditBench 数据集拉图（后端 requests）"""
+    rel_path = rel_path.lstrip("/")
+    if rel_path in _image_cache:
+        return _image_cache[rel_path]
+    url = f"{BASE_HF_URL}/{rel_path}"
+    resp = get_url_response(url)
+    if not resp:
+        return None
+    img = Image.open(BytesIO(resp.content)).convert("RGB")
+    _image_cache[rel_path] = img
+    return img
+
+def pil_to_base64(img):
+    if img is None:
+        return ""
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    s = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{s}"
 
 ROOT_DIR = Path(__file__).resolve().parent
 
 # =============== data.json 还是放在 Space 本地 =================
+# 如果想把 data.json 也放到 dataset 里，可以改成用 hf_hub_download（后面我再写）
 JSON_PATH = ROOT_DIR / "data.json"
 
 # =============== 读 json & 建索引 =================
@@ -21,7 +60,8 @@ ALL_NAMES = sorted({item["name"] for item in data})
 # (name, idx_str) -> item
 INDEX_MAP = {(item["name"], str(item["idx"])): item for item in data}
 
-# =============== baseline 模型列表 =================
+# =============== baseline 模型列表（现在不能再 os.listdir，本地没图了） =================
+# 建议直接手动写齐你有的模型
 ALL_MODELS = [
     "Bagel-Think",
     "DreamOmni2",
@@ -58,20 +98,27 @@ def get_indices_for_name(name: str):
 
 
 def render_img_html(rel_path: str, max_h=512):
-    """
-    像第二段代码那样，直接返回远程 URL，不在后端拉图 / 编码。
-    rel_path 例如: 'original_image/jewel2/0001.png'
-    """
     if not rel_path:
         return "<p>No original image.</p>"
 
-    rel_path = str(rel_path).lstrip("/")  # 防止开头有 '/'
-    src = f"{HF_BASE_URL}/{rel_path}"
-    return f'<img src="{src}" style="max-width:100%; max-height:{max_h}px;">'
+    key = (rel_path, None)  # 不区分高度
+    if key in _html_cache:
+        data_url = _html_cache[key]
+        return f'<img src="{data_url}" style="max-width:100%; max-height:{max_h}px;">'
+
+    img = load_image_uniredit(rel_path)
+    if img is None:
+        return "<p>Failed to load image.</p>"
+
+    # 不再 resize，直接按原分辨率编码
+    data_url = pil_to_base64(img)
+    _html_cache[key] = data_url
+
+    return f'<img src="{data_url}" style="max-width:100%; max-height:{max_h}px;">'
 
 
 def get_baseline_gallery(name: str, idx: str, models):
-    """生成 baseline 图像的 HTML 表格（和第二段类似，img src 直接是 URL）."""
+    """生成 baseline 图像的 HTML 表格（使用远程 URL）."""
     if not name or not idx:
         return "<p>Please select name and idx.</p>"
 
@@ -90,7 +137,9 @@ def get_baseline_gallery(name: str, idx: str, models):
         # 第一行：模型名
         html += "<tr>"
         for m in sub_models:
-            html += f'<td width="{WIDTH}%" style="text-align:center;"><h4>{m}</h4></td>'
+            html += (
+                f'<td width="{WIDTH}%" style="text-align:center;"><h4>{m}</h4></td>'
+            )
         for _ in range(N_COL - len(sub_models)):
             html += f'<td width="{WIDTH}%"></td>'
         html += "</tr>"
@@ -98,10 +147,8 @@ def get_baseline_gallery(name: str, idx: str, models):
         # 第二行：对应图片
         html += "<tr>"
         for m in sub_models:
-            # 对应 dataset 目录：Unireditbench_baseline_images/{model}/{name}/{idx}.png
             rel_path = f"Unireditbench_baseline_images/{m}/{name}/{idx}.png"
-            src = f"{HF_BASE_URL}/{rel_path}"
-            cell = f'<img src="{src}" style="max-width:100%; max-height:256px;">'
+            cell = render_img_html(rel_path, max_h=256)
             html += f'<td width="{WIDTH}%" style="text-align:center;">{cell}</td>'
         for _ in range(N_COL - len(sub_models)):
             html += f'<td width="{WIDTH}%"></td>'
@@ -133,7 +180,7 @@ def load_sample(name, idx, selected_models):
     instruction = item.get("instruction", "")
     rules = item.get("rules", "")
 
-    # data.json 里 original_image_path 仍然保持为相对路径：
+    # data.json 里原本的 original_image_path 建议直接保持为相对路径：
     #   "original_image_path": "original_image/jewel2/0001.png"
     orig_rel = item.get("original_image_path", "")
     orig_html = render_img_html(orig_rel, max_h=512)
@@ -142,7 +189,6 @@ def load_sample(name, idx, selected_models):
 
     return info_md, instruction, rules, orig_html, gallery_html
 
-
 def _step_idx(name: str, idx: str, direction: int):
     """direction=-1 上一个, direction=+1 下一个；到边界就停住。"""
     idxs = get_indices_for_name(name)
@@ -206,11 +252,12 @@ with gr.Blocks() as demo:
                 prev_button = gr.Button("Prev")
                 next_button = gr.Button("Next")
                 load_button = gr.Button("Load sample")
+                
 
             model_checkboxes = gr.CheckboxGroup(
                 label="Baselines to show",
-                choices=BASELINE_MODELS,
-                value=BASELINE_MODELS,
+                choices=BASELINE_MODELS,     # 还是所有模型都可以选
+                value=PRIORITY_MODELS,       # ✅ 默认只勾选优先那几个
             )
 
         with gr.Column(scale=2):
@@ -226,6 +273,7 @@ with gr.Blocks() as demo:
             gr.Markdown("### Original Image")
             orig_image_html = gr.HTML()
 
+
     gallery_html = gr.HTML()
 
     name_dropdown.change(
@@ -239,7 +287,7 @@ with gr.Blocks() as demo:
         inputs=[name_dropdown, idx_dropdown, model_checkboxes],
         outputs=[info_markdown, instruction_box, rules_box, orig_image_html, gallery_html]
     )
-
+    
     # Prev / Next：会同时更新 idx_dropdown + 刷新所有内容
     prev_button.click(
         fn=prev_sample,
@@ -255,4 +303,5 @@ with gr.Blocks() as demo:
 
 
 if __name__ == "__main__":
+    # 不再需要 allowed_paths / set_static_paths
     demo.launch(server_name="0.0.0.0", server_port=7860)