Update app.py

app.py

@@ -2,10 +2,46 @@ import os
 import json
 import gradio as gr
 from pathlib import Path
+from io import BytesIO
+import base64
+import requests
 
 # =============== 配置：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
 
@@ -60,16 +96,22 @@ def get_indices_for_name(name: str):
 
 
 def render_img_html(rel_path: str, max_h=512):
-    """
-    现在不再检查本地文件，直接把相对路径拼到 HF dataset 的 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;">'
+    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;">'
 
 
 def get_baseline_gallery(name: str, idx: str, models):
@@ -102,14 +144,9 @@ def get_baseline_gallery(name: str, idx: str, models):
         # 第二行：对应图片
         html += "<tr>"
         for m in sub_models:
-            # 关键：这里直接拼 dataset 里的路径
-            # 对应 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;">'
-            html += (
-                f'<td width="{WIDTH}%" style="text-align:center;">{cell}</td>'
-            )
+            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>'
         html += "</tr>"