Update app_v2.py

app_v2.py

@@ -1,18 +1,18 @@
 """
 ==========================================
-Image Anonymizer 
+Screenshot Anonymizer — v2 (tool revision)
 ==========================================
 """
 
 import base64
 import functools
 import io
-import json
+import traceback
 from pathlib import Path
 
 import gradio as gr
-from fastapi import File, UploadFile
 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, Response
+from gradio.data_classes import FileData
 from PIL import Image
 
 from app import (
@@ -56,46 +56,11 @@ async def homepage():
     return FRONTEND_HTML
 
 
-@server.post("/api/detect")
-async def detect(file: UploadFile = File(...)):
-    suffix = Path(file.filename or "").suffix.lower()
-    if suffix not in (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"):
-        return JSONResponse({"error": f"Unsupported image type: {suffix or '(none)'}"}, 400)
-    try:
-        img_bytes = await file.read()
-        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-    except Exception as e:
-        return JSONResponse({"error": f"Could not read image: {e}"}, 400)
-
-    ocr = ocr_image(img)
-    if not ocr["text"].strip():
-        return JSONResponse({"error": "No text detected in the image."}, 400)
-
-    try:
-        source_text, spans = run_pii_analysis(ocr["text"])
-    except Exception as e:
-        return JSONResponse({"error": f"PII analysis failed: {e}"}, 500)
-
-    if source_text != ocr["text"]:
-        spans = [s for s in spans if s["end"] <= len(ocr["text"])]
-
-    boxes = map_spans_to_boxes(ocr["words"], spans)
-
-    buf = io.BytesIO(); img.save(buf, format="PNG")
-    data_url = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
-
-    return JSONResponse({
-        "filename": file.filename,
-        "image": data_url,
-        "width": img.width, "height": img.height,
-        "boxes": boxes,
-        "text": ocr["text"],
-        "spans": spans,
-        "categories_meta": {k: {"color": v["color"], "label": v["label"]}
-                            for k, v in CATEGORIES_META.items()},
-    })
-
-
+# The /api/examples and /examples/{name} routes below serve static
+# example thumbnails and originals from disk. They are plain FastAPI
+# routes because they do no GPU / queued compute — they just read
+# files, which is exactly the pattern the gradio.Server blog
+# recommends plain @server.get for.
 @server.get("/api/examples")
 async def api_examples():
     return JSONResponse({"examples": _list_examples()})
@@ -118,37 +83,51 @@ async def get_example(name: str, thumb: int = 0):
     return FileResponse(path, headers={"Cache-Control": "public, max-age=3600"})
 
 
+# The compute endpoint: goes through Gradio's queue, plays nicely with
+# @spaces.GPU on ZeroGPU, and is callable by both the browser via the
+# @gradio/client JS client AND by Python users via gradio_client.
 @server.api(name="anonymize_screenshot")
-def anonymize_screenshot_api(image_path: str) -> str:
-    """Gradio API: accepts an image (via gradio_client.handle_file) and
-    returns detected redaction boxes as JSON.
+def anonymize_screenshot_api(image: FileData) -> dict:
+    """OCR + PII-filter an uploaded image.
 
-    Tolerant of both a plain path string and a FileData dict, since
-    different gradio_client versions pass each.
+    Input: FileData from `handle_file(file)` (JS client) or
+    `gradio_client.handle_file(path)` (Python client).
+    Output: dict with the base-64 image, OCR text, detected spans,
+    per-span pixel boxes, and the category color/label table.
     """
-    import traceback
     try:
-        if isinstance(image_path, dict):
-            image_path = image_path.get("path") or image_path.get("url") or ""
-        if not image_path or not isinstance(image_path, str):
-            return json.dumps({"error": f"expected image path, got {type(image_path).__name__}"})
+        path = image.get("path") or image.get("url") or ""
+        if not path:
+            return {"error": "expected an image file"}
 
-        img = Image.open(image_path).convert("RGB")
+        img = Image.open(path).convert("RGB")
         ocr = ocr_image(img)
         if not ocr["text"].strip():
-            return json.dumps({
-                "width": img.width, "height": img.height,
-                "boxes": [], "text": "", "spans": [],
-            })
-        _, spans = run_pii_analysis(ocr["text"])
+            return {"error": "No text detected in the image."}
+
+        source_text, spans = run_pii_analysis(ocr["text"])
+        if source_text != ocr["text"]:
+            spans = [s for s in spans if s["end"] <= len(ocr["text"])]
         boxes = map_spans_to_boxes(ocr["words"], spans)
-        return json.dumps({
+
+        buf = io.BytesIO(); img.save(buf, format="PNG")
+        data_url = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+        return {
+            "filename": Path(path).name,
+            "image": data_url,
             "width": img.width, "height": img.height,
-            "boxes": boxes, "text": ocr["text"], "spans": spans,
-        }, ensure_ascii=False)
+            "boxes": boxes,
+            "text": ocr["text"],
+            "spans": spans,
+            "categories_meta": {
+                k: {"color": v["color"], "label": v["label"]}
+                for k, v in CATEGORIES_META.items()
+            },
+        }
     except Exception as e:
         traceback.print_exc()
-        return json.dumps({"error": f"{type(e).__name__}: {e}"})
+        return {"error": f"{type(e).__name__}: {e}"}
 
 
 # =====================================================================
@@ -160,7 +139,7 @@ FRONTEND_HTML = r"""<!DOCTYPE html>
 <head>
 <meta charset="UTF-8">
 <meta name="viewport" content="width=device-width,initial-scale=1">
-<title>Image Anonymizer</title>
+<title>Screenshot Anonymizer</title>
 <link rel="preconnect" href="https://fonts.googleapis.com">
 <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
 <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500&family=Lora:ital,wght@0,400;0,500;1,400&display=swap" rel="stylesheet">
@@ -316,7 +295,7 @@ svg{display:block;flex-shrink:0}
 
 .landing-foot{
   max-width:720px;margin:.5rem auto 2rem;padding:0 1.5rem;
-  font-family:var(--mono);font-size:12px;color:var(--text3);line-height:1.7;
+  font-family:var(--mono);font-size:10.5px;color:var(--text3);line-height:1.7;
 }
 .landing-foot a{color:var(--text2);text-decoration:none;border-bottom:.5px dotted var(--text3)}
 .landing-foot a:hover{color:var(--text)}
@@ -543,7 +522,7 @@ svg{display:block;flex-shrink:0}
         <rect x="2" y="9"  width="16" height="3" rx="0.5" fill="currentColor"/>
         <rect x="2" y="14" width="8"  height="3" rx="0.5" fill="currentColor"/>
       </svg>
-      <span class="wordmark">Image Anonymizer</span>
+      <span class="wordmark">Screenshot Anonymizer</span>
       <span class="version">v0.3 · beta</span>
     </div>
     <div class="spacer"></div>
@@ -551,8 +530,8 @@ svg{display:block;flex-shrink:0}
   </header>
 
   <div class="landing-content">
-    <h1 class="headline">Redact images before you share them.</h1>
-    <p class="subtitle">First OCR finds the text. Then OpenAI's Privacy Filter marks names, emails, phones, addresses, dates, URLs, accounts, and secrets. You approve what gets hidden before it ever leaves the page.</p>
+    <h1 class="headline">Redact screenshots before you post them.</h1>
+    <p class="subtitle">OCR finds the text. The privacy filter marks names, emails, phones, addresses, dates, URLs, accounts, and secrets. You approve what gets hidden before it ever leaves the page.</p>
 
     <label class="dropzone" id="dropzone">
       <input type="file" id="file-input" accept="image/png,image/jpeg,image/webp,image/bmp,image/tiff">
@@ -561,13 +540,13 @@ svg{display:block;flex-shrink:0}
         <rect x="2" y="12" width="38" height="5" rx="1" fill="currentColor"/>
         <rect x="2" y="21" width="18" height="5" rx="1" fill="currentColor"/>
       </svg>
-      <div class="dz-text">Drop an image, paste from clipboard, or click to browse</div>
+      <div class="dz-text">Drop a screenshot, paste from clipboard, or click to browse</div>
       <div class="dz-hint">png · jpg · webp · bmp · tiff</div>
     </label>
 
     <div class="example-wrap">
       <div class="example-head">
-        <span class="title">try an example · click to load (wait for a few seconds for the example to load)</span>
+        <span class="title">try an example · click to load</span>
         <span class="note">all content is fictitious · mock data for testing only</span>
       </div>
       <div class="example-scroll" id="example-scroll">
@@ -577,8 +556,8 @@ svg{display:block;flex-shrink:0}
   </div>
 
   <div class="landing-foot">
-    Model: <b>OpenAI Privacy Filter</b> · 1.5b params, 50m active · Apache 2.0<br>
-    OCR: Tesseract 5 · <b>Processing: server-side using gr.Server, ZEROGPU</b> · edits stay in your browser
+    pii: openai/privacy-filter · 1.5b params, 50m active · apache 2.0<br>
+    ocr: rednote-hilab/dots.ocr · 3b vlm, top-3 on olmocr-bench · edits stay in your browser
   </div>
 </div>
 
@@ -591,11 +570,11 @@ svg{display:block;flex-shrink:0}
         <rect x="2" y="9"  width="16" height="3" rx="0.5" fill="currentColor"/>
         <rect x="2" y="14" width="8"  height="3" rx="0.5" fill="currentColor"/>
       </svg>
-      <span class="wordmark">Image Anonymizer</span>
+      <span class="wordmark">Screenshot Anonymizer</span>
     </div>
     <span class="version" id="meta-info">—</span>
     <div class="spacer"></div>
-    <button class="btn-link" onclick="resetView()">new image</button>
+    <button class="btn-link" onclick="resetView()">new screenshot</button>
   </header>
 
   <div class="error-banner" id="error-banner"></div>
@@ -689,7 +668,17 @@ svg{display:block;flex-shrink:0}
 <div id="loading"><div class="spinner"></div><p>ocr → privacy filter → map to pixels</p></div>
 <div class="toast" id="toast"></div>
 
-<script>
+<script type="module">
+// ═══════════════════════════════════���══════════════════════════════
+// Gradio JS client — talks to the queued @server.api routes so that
+// requests are serialized, progress is tracked, and ZeroGPU's
+// @spaces.GPU allocator gets invoked correctly. A plain fetch() to a
+// FastAPI route would bypass all of that.
+// ══════════════════════════════════════════════════════════════════
+import { Client, handle_file } from "https://cdn.jsdelivr.net/npm/@gradio/client/dist/index.min.js";
+
+const clientPromise = Client.connect(window.location.origin);
+
 // ══════════════════════════════════════════════════════════════════
 // State
 // ══════════════════════════════════════════════════════════════════
@@ -778,14 +767,20 @@ async function uploadFile(file) {
   if (!file.type || !file.type.startsWith('image/')) { showError('please drop an image file.'); return; }
   document.getElementById('loading').style.display = 'flex';
   document.getElementById('landing').style.display = 'none';
-  const form = new FormData(); form.append('file', file);
   try {
-    const r = await fetch('/api/detect', { method:'POST', body: form });
-    const d = await r.json();
+    const client = await clientPromise;
+    const result = await client.predict("/anonymize_screenshot", {
+      image: handle_file(file),
+    });
+    // @server.api returns a dict; the client wraps outputs in result.data[]
+    const d = result.data[0] || {};
     if (d.error) { showError(d.error); return; }
+    // The server returns Path(path).name as `filename`; fall back to the
+    // original File.name for display continuity.
+    if (!d.filename) d.filename = file.name;
     await initEditor(d);
   } catch (e) {
-    showError('analysis failed: ' + e.message);
+    showError('analysis failed: ' + (e && e.message ? e.message : e));
   } finally {
     document.getElementById('loading').style.display = 'none';
   }
@@ -1291,6 +1286,14 @@ function toast(msg) {
   clearTimeout(toastTimer);
   toastTimer = setTimeout(() => t.classList.remove('show'), 2000);
 }
+
+// Module scripts don't expose top-level names to the global scope, but
+// several buttons in the HTML use inline onclick="foo()" handlers —
+// bridge them explicitly.
+Object.assign(window, {
+  resetView, zoomStep, zoomFit, zoomReset,
+  downloadImage, copyToClipboard, exportText,
+});
 </script>
 </body>
 </html>"""