Spaces:

stephenebert
/

Image_Tagger

Sleeping

App Files Files Community

stephenebert commited on Aug 13, 2025

Commit

cb8b918

verified ·

1 Parent(s): a6de95a

Update app.py

Browse files

Files changed (1) hide show

app.py +73 -27

app.py CHANGED Viewed

@@ -12,10 +12,10 @@ from fastapi.responses import HTMLResponse, JSONResponse
 from PIL import Image
 import gradio as gr
-# Uses your existing tagger.py (BLIP loaded there; returns top-k tags and writes sidecar)
 import tagger as tg
-# ---------- FastAPI base app ----------
 app = FastAPI(title="Image Tagger API")
 app.add_middleware(
@@ -27,31 +27,33 @@ DATA_DIR = Path("/app/data")
 DATA_DIR.mkdir(parents=True, exist_ok=True)
-def _caption_with_tagger(img: Image.Image) -> str:
-    """Optional caption via BLIP objects already loaded in tagger.py."""
     try:
         proc = getattr(tg, "_processor")
         model = getattr(tg, "_model")
         ids = model.generate(**proc(images=img, return_tensors="pt"), max_length=30)
         return proc.decode(ids[0], skip_special_tokens=True)
     except Exception:
         return ""
-# ---------- API endpoints ----------
 @app.get("/healthz")
 def healthz():
     return {"ok": True}
 @app.get("/", response_class=HTMLResponse)
 def root():
     return """<!doctype html>
 <html>
   <head><meta charset="utf-8" /><title>Image Tagger API</title></head>
-  <body style="font-family: system-ui; max-width: 720px; margin: 40px auto">
     <h2>Image Tagger API</h2>
-    <p>Try the Swagger docs at <a href="/docs">/docs</a> or the UI at <a href="/ui">/ui</a>.</p>
-    <form action="/upload" method="post" enctype="multipart/form-data" style="display:grid; gap:12px">
       <input type="file" name="file" accept="image/png,image/jpeg,image/webp" required />
       <label>top_k: <input type="number" name="top_k" value="5" min="1" max="20" /></label>
       <button type="submit">Upload</button>
@@ -59,6 +61,7 @@ def root():
   </body>
 </html>"""
 @app.post("/upload")
 async def upload_image(
     file: UploadFile = File(...),
@@ -68,35 +71,78 @@ async def upload_image(
     img = Image.open(io.BytesIO(content)).convert("RGB")
     stem = Path(file.filename).stem
-    tags: List[str] = tg.tag_pil_image(img, stem, top_k=top_k)  # tagger writes sidecar in /app/data (per your tagger)
-    caption = _caption_with_tagger(img)
     payload = {"filename": file.filename, "caption": caption, "tags": tags}
-    # Save a copy for convenience
     (DATA_DIR / f"{stem}.json").write_text(json.dumps(payload, indent=2))
     img.save(DATA_DIR / file.filename)
     return JSONResponse(payload)
-# ---------- Gradio UI at /ui ----------
-def _gr_predict(img: Image.Image, k: int):
     if img is None:
-        return "", ""
-    tags = tg.tag_pil_image(img.convert("RGB"), "ui_upload", top_k=int(k))
-    caption = _caption_with_tagger(img)
-    return caption, ", ".join(tags)
-with gr.Blocks(css=".gradio-container {max-width: 860px !important}") as demo:
-    gr.Markdown("## Image Tagger — UI")
     with gr.Row():
-        in_img = gr.Image(type="pil", label="Image")
-        k = gr.Slider(1, 20, value=5, step=1, label="Top-k tags")
-    run = gr.Button("Tag image")
-    out_caption = gr.Textbox(label="Caption", lines=2)
-    out_tags = gr.Textbox(label="Tags (comma-separated)", lines=2)
-    run.click(_gr_predict, inputs=[in_img, k], outputs=[out_caption, out_tags])
-# IMPORTANT: export lowercase `app` (Uvicorn expects app:app)
 app = gr.mount_gradio_app(app, demo, path="/ui")

 from PIL import Image
 import gradio as gr
+# Your BLIP + tagging logic lives here
 import tagger as tg
+# -------------------- FastAPI --------------------
 app = FastAPI(title="Image Tagger API")
 app.add_middleware(
 DATA_DIR.mkdir(parents=True, exist_ok=True)
+def _caption_with_blip(img: Image.Image) -> str:
+    """Produce a caption using the BLIP objects that tagger.py already loaded."""
     try:
         proc = getattr(tg, "_processor")
         model = getattr(tg, "_model")
         ids = model.generate(**proc(images=img, return_tensors="pt"), max_length=30)
         return proc.decode(ids[0], skip_special_tokens=True)
     except Exception:
+        # If anything goes wrong, don't break the UI — just return empty caption
         return ""
 @app.get("/healthz")
 def healthz():
     return {"ok": True}
 @app.get("/", response_class=HTMLResponse)
 def root():
+    # tiny landing page for folks who hit the root
     return """<!doctype html>
 <html>
   <head><meta charset="utf-8" /><title>Image Tagger API</title></head>
+  <body style="font-family:system-ui;max-width:760px;margin:40px auto">
     <h2>Image Tagger API</h2>
+    <p>Try the UI at <a href="/ui">/ui</a> or the API docs at <a href="/docs">/docs</a>.</p>
+    <form action="/upload" method="post" enctype="multipart/form-data" style="display:grid;gap:12px">
       <input type="file" name="file" accept="image/png,image/jpeg,image/webp" required />
       <label>top_k: <input type="number" name="top_k" value="5" min="1" max="20" /></label>
       <button type="submit">Upload</button>
   </body>
 </html>"""
 @app.post("/upload")
 async def upload_image(
     file: UploadFile = File(...),
     img = Image.open(io.BytesIO(content)).convert("RGB")
     stem = Path(file.filename).stem
+    # Use your tagger (returns top-k tags; no POS filters)
+    tags: List[str] = tg.tag_pil_image(img, stem, top_k=top_k)
+    caption = _caption_with_blip(img)
     payload = {"filename": file.filename, "caption": caption, "tags": tags}
+    # Persist a sidecar + the uploaded image (handy for debugging)
     (DATA_DIR / f"{stem}.json").write_text(json.dumps(payload, indent=2))
     img.save(DATA_DIR / file.filename)
     return JSONResponse(payload)
+# -------------------- Gradio UI (/ui) --------------------
+def _predict_ui(img: Image.Image | None, k: int):
     if img is None:
+        return "", []  # caption, table rows
+    img = img.convert("RGB")
+    caption = _caption_with_blip(img)
+    tags = tg.tag_pil_image(img, "ui_upload", top_k=int(k))
+    # Build a ranked table: [["#1", "lion"], ["#2", "rock"], ...]
+    rows = [[f"{i+1}", tag] for i, tag in enumerate(tags)]
+    return caption, rows
+CSS = """
+.gradio-container { max-width: 1200px !important; }
+.caption-box textarea { font-size: 16px !important; }
+.rank-table thead th { font-weight: 600; }
+"""
+with gr.Blocks(css=CSS, theme=gr.themes.Default(primary_hue="blue")) as demo:
+    gr.Markdown("### 🔎 Image Tagger (BLIP ➜ Caption + Tags)")
+    with gr.Row(equal_height=True):
+        # LEFT: big image
+        with gr.Column(scale=7):
+            in_img = gr.Image(type="pil", label="Upload Image", height=540)
+        # RIGHT: caption + ranked list
+        with gr.Column(scale=5):
+            gr.Markdown("**BLIP Generated Caption**")
+            out_caption = gr.Textbox(
+                label=None, lines=2, elem_classes=["caption-box"], interactive=False
+            )
+            gr.Markdown("**Top-k Tags (ranked)**")
+            out_table = gr.Dataframe(
+                headers=["Rank", "Tag"],
+                datatype=["str", "str"],
+                col_count=(2, "fixed"),
+                row_count=(0, "dynamic"),
+                wrap=True,
+                interactive=False,
+                elem_classes=["rank-table"],
+                height=380,
+            )
     with gr.Row():
+        k = gr.Slider(1, 20, value=5, step=1, label="Number of Tags (k)")
+    with gr.Row():
+        clear_btn = gr.ClearButton(
+            [in_img, out_caption, out_table],
+            value="Clear",
+            size="sm",
+        )
+        submit_btn = gr.Button("Submit", variant="primary")
+    submit_btn.click(_predict_ui, inputs=[in_img, k], outputs=[out_caption, out_table])
+# Export lowercase `app` for Uvicorn ("app:app")
 app = gr.mount_gradio_app(app, demo, path="/ui")