Update app.py

app.py

@@ -1,38 +1,70 @@
-
-import io, math
+import io
 from PIL import Image, ImageChops, ImageStat, ExifTags
 import gradio as gr
 
+# 可能出现在 EXIF 里的生成器关键词（可自行扩展）
 GENERATOR_KEYWORDS = [
     "stable diffusion", "stability.ai", "sdxl", "midjourney", "dall", "openai",
     "novelai", "leonardo", "kaiber", "flux", "comfyui", "automatic1111", "invokeai"
 ]
 
-def to_rgb(img):
-    if img.mode in ["RGBA", "P"]:
-        return img.convert("RGB")
-    if img.mode != "RGB":
-        return img.convert("RGB")
-    return img
+def to_rgb_flat(img, bg=(255, 255, 255)):
+    """确保得到 RGB；遇到 RGBA/带透明通道时做白底合成，避免 JPEG 保存时报错。"""
+    if img.mode == "RGB":
+        return img
+    if img.mode in ("RGBA", "LA", "P"):
+        bg_img = Image.new("RGB", img.size, bg)
+        if img.mode == "P":
+            img = img.convert("RGBA")
+        bg_img.paste(img, mask=img.split()[-1] if "A" in img.getbands() else None)
+        return bg_img
+    return img.convert("RGB")
+
+def resize_max(img, max_side=1024):
+    """把最长边限制到 1024，降低内存占用并避免某些编码错误。"""
+    w, h = img.size
+    m = max(w, h)
+    if m <= max_side:
+        return img
+    scale = max_side / float(m)
+    return img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
 
 def compute_ela_score(img, quality=95):
-    # Error Level Analysis: recompress at given JPEG quality, then compute difference stats
-    img_rgb = to_rgb(img)
-    buf = io.BytesIO()
-    img_rgb.save(buf, "JPEG", quality=quality, optimize=True)
-    buf.seek(0)
-    recompressed = Image.open(buf).convert("RGB")
-    ela = ImageChops.difference(img_rgb, recompressed)
-    stat = ImageStat.Stat(ela)
-    # mean and std over RGB channels
-    mean = sum(stat.mean) / len(stat.mean)
-    std = sum(stat.stddev) / len(stat.stddev)
-    # normalize to 0-1 rough scale (not used for decision, but could be returned)
-    mean_norm = min(mean / 10.0, 1.0)
-    std_norm  = min(std / 10.0, 1.0)
-    return mean, std, mean_norm, std_norm
+    """
+    ELA（误差层分析）：以给定 JPEG 质量重压一遍，然后计算差异图的均值/方差。
+    若失败（例如非 JPEG 友好的模式/编解码异常），返回 (None, None) 并由上层降级处理。
+    """
+    try:
+        img_rgb = to_rgb_flat(img)
+        img_rgb = resize_max(img_rgb, 1024)
+        buf = io.BytesIO()
+        img_rgb.save(buf, "JPEG", quality=quality, optimize=True)
+        buf.seek(0)
+        recompressed = Image.open(buf).convert("RGB")
+        ela = ImageChops.difference(img_rgb, recompressed)
+        stat = ImageStat.Stat(ela)
+        mean = float(sum(stat.mean) / len(stat.mean))
+        std = float(sum(stat.stddev) / len(stat.stddev))
+        return mean, std
+    except Exception:
+        # 再尝试一次更保守的质量设置
+        try:
+            img_rgb = to_rgb_flat(img)
+            img_rgb = resize_max(img_rgb, 1024)
+            buf = io.BytesIO()
+            img_rgb.save(buf, "JPEG", quality=85)
+            buf.seek(0)
+            recompressed = Image.open(buf).convert("RGB")
+            ela = ImageChops.difference(img_rgb, recompressed)
+            stat = ImageStat.Stat(ela)
+            mean = float(sum(stat.mean) / len(stat.mean))
+            std = float(sum(stat.stddev) / len(stat.stddev))
+            return mean, std
+        except Exception:
+            return None, None  # 彻底放弃 ELA，交由上层“降级”
 
 def extract_exif_flags(img):
+    """读取少量常见 EXIF 字段，并搜寻生成器关键词。异常直接吞掉，返回尽量多的信息。"""
     exif = {}
     try:
         raw = img.getexif()
@@ -44,69 +76,79 @@ def extract_exif_flags(img):
 
     exif_str = " ".join([str(v).lower() for v in exif.values()]) if exif else ""
     has_camera_fields = any(tag in exif for tag in ["Make", "Model", "LensModel", "DateTimeOriginal"])
-    has_software = "Software" in exif
     has_generator_kw = any(kw in exif_str for kw in GENERATOR_KEYWORDS)
     empty_exif = (len(exif) == 0)
 
     preview = {}
-    for k in ["Make","Model","LensModel","Software","DateTimeOriginal"]:
+    for k in ["Make", "Model", "LensModel", "Software", "DateTimeOriginal"]:
         if k in exif:
             preview[k] = str(exif[k])
 
     return {
         "has_camera_fields": has_camera_fields,
-        "has_software": has_software,
         "has_generator_kw": has_generator_kw,
         "empty_exif": empty_exif,
         "exif_preview": preview
     }
 
 def ai_likelihood(img):
-    if img is None:
-        return {"error": "No image uploaded."}
-    info = extract_exif_flags(img)
-    mean, std, mean_norm, std_norm = compute_ela_score(img)
-
-    score = 0.5
-    reasons = []
-
-    if info["has_generator_kw"]:
-        score += 0.4
-        reasons.append("Metadata contains generator keywords (e.g., Stable Diffusion/Midjourney).")
-    if info["has_camera_fields"]:
-        score -= 0.2
-        reasons.append("Camera EXIF fields found (Make/Model/Lens), suggesting real capture.")
-    if info["empty_exif"]:
-        score += 0.1
-        reasons.append("No EXIF found (common in exported AI images or screenshots).")
-    # ELA heuristic: AI images often show uniform reconstruction errors
-    if mean < 2.0 and std < 2.0:
-        score += 0.15
-        reasons.append("ELA mean/std are very low → uniform compression error (AI-like).")
-    elif mean > 4.0 or std > 4.0:
-        score -= 0.05
-        reasons.append("ELA mean/std are higher → natural camera/post-processing artifacts (Real-like).")
-
-    score = max(0.0, min(1.0, score))
-    label = "Likely AI" if score >= 0.6 else ("Uncertain" if 0.4 <= score < 0.6 else "Likely Real")
-
-    explanation = {
-        "label": label,
-        "ai_probability": round(score, 3),
-        "ela_mean": round(mean, 3),
-        "ela_std": round(std, 3),
-        "exif": info["exif_preview"],
-        "notes": reasons or ["No strong signals; result uncertain."]
-    }
-    return explanation
+    """
+    主入口：任何异常都捕获，返回 JSON 友好信息而不是让前端报“错误”。
+    """
+    try:
+        if img is None:
+            return {"label": "Error", "message": "No image uploaded."}
+
+        info = extract_exif_flags(img)
+        ela_mean, ela_std = compute_ela_score(img)
+
+        # 初始分数（0.5 = 不确定）
+        score = 0.5
+        reasons = []
+
+        if info["has_generator_kw"]:
+            score += 0.4
+            reasons.append("Metadata contains generator keywords (e.g., Stable Diffusion/Midjourney).")
+        if info["has_camera_fields"]:
+            score -= 0.2
+            reasons.append("Camera EXIF fields found (Make/Model/Lens/DateTimeOriginal).")
+        if info["empty_exif"]:
+            score += 0.1
+            reasons.append("No EXIF found (common in exported AI images or screenshots).")
+
+        if ela_mean is not None and ela_std is not None:
+            if ela_mean < 2.0 and ela_std < 2.0:
+                score += 0.15
+                reasons.append("ELA mean/std are very low → uniform compression error (AI-like).")
+            elif ela_mean > 4.0 or ela_std > 4.0:
+                score -= 0.05
+                reasons.append("ELA mean/std are higher → natural camera/post-processing artifacts (Real-like).")
+        else:
+            reasons.append("ELA failed (unsupported format/codec); decision based on metadata only.")
+
+        score = max(0.0, min(1.0, score))
+        label = "Likely AI" if score >= 0.6 else ("Uncertain" if 0.4 <= score < 0.6 else "Likely Real")
+
+        return {
+            "label": label,
+            "ai_probability": round(score, 3),
+            "ela_mean": None if ela_mean is None else round(ela_mean, 3),
+            "ela_std": None if ela_std is None else round(ela_std, 3),
+            "exif": info["exif_preview"],
+            "notes": reasons or ["No strong signals; result uncertain."]
+        }
+
+    except Exception as e:
+        # 兜底：把异常显示在 JSON 里，便于你在前端看到具体原因
+        return {"label": "Error", "message": str(e)}
 
 with gr.Blocks() as demo:
     gr.Markdown("""
     # 🕵️ FakeSpotter (Heuristic Demo)
     Upload an image to estimate whether it is **AI-generated** or **Real** using simple FREE heuristics:
     - Metadata scan (generator keywords vs. camera EXIF)
-    - ELA (Error Level Analysis) statistics
-    > ⚠️ This is a classroom demo, **not a forensic tool**.
+    - ELA (Error Level Analysis) statistics  
+    > ⚠️ Classroom demo, **not** a forensic tool.
     """)
     inp = gr.Image(type="pil", label="Upload image")
     out = gr.JSON(label="Result")