Update app.py

app.py

@@ -2,52 +2,14 @@ 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 = {}
-
-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 & 建索引 =================
@@ -59,8 +21,7 @@ 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 模型列表（现在不能再 os.listdir，本地没图了） =================
-# 建议直接手动写齐你有的模型
+# =============== baseline 模型列表 =================
 ALL_MODELS = [
     "Bagel-Think",
     "DreamOmni2",
@@ -97,26 +58,20 @@ 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>"
 
-    img = load_image_uniredit(rel_path)
-    if img is None:
-        return "<p>Failed to load image.</p>"
-
-    # 可选：缩放一下高度
-    h = img.height
-    w = img.width
-    if h > max_h:
-        ratio = max_h / h
-        img = img.resize((int(w * ratio), max_h), Image.Resampling.LANCZOS)
-
-    data_url = pil_to_base64(img)
-    return f'<img src="{data_url}" style="max-width:100%; max-height:{max_h}px;">'
+    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;">'
 
 
 def get_baseline_gallery(name: str, idx: str, models):
-    """生成 baseline 图像的 HTML 表格（使用远程 URL）."""
+    """生成 baseline 图像的 HTML 表格（和第二段类似，img src 直接是 URL）."""
     if not name or not idx:
         return "<p>Please select name and idx.</p>"
 
@@ -135,9 +90,7 @@ 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>"
@@ -145,8 +98,10 @@ 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"
-            cell = render_img_html(rel_path, max_h=256)
+            src = f"{HF_BASE_URL}/{rel_path}"
+            cell = f'<img src="{src}" style="max-width:100%; max-height:256px;">'
             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>'
@@ -178,7 +133,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)
@@ -187,6 +142,7 @@ 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)
@@ -250,7 +206,6 @@ 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",
@@ -271,7 +226,6 @@ with gr.Blocks() as demo:
             gr.Markdown("### Original Image")
             orig_image_html = gr.HTML()
 
-
     gallery_html = gr.HTML()
 
     name_dropdown.change(
@@ -285,7 +239,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,
@@ -301,5 +255,4 @@ with gr.Blocks() as demo:
 
 
 if __name__ == "__main__":
-    # 不再需要 allowed_paths / set_static_paths
     demo.launch(server_name="0.0.0.0", server_port=7860)