Spaces:

stephenebert
/

Image_Tagger

Sleeping

App Files Files Community

stephenebert commited on Aug 13, 2025

Commit

d3e85eb

verified ·

1 Parent(s): 0e857c8

Update app.py

Browse files

Files changed (1) hide show

app.py +90 -106

app.py CHANGED Viewed

@@ -1,42 +1,36 @@
-# app.py
 from __future__ import annotations
 import io
-import json
 from pathlib import Path
 from typing import List
-from fastapi import FastAPI, UploadFile, File, Query
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse, JSONResponse
-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(
-    CORSMiddleware,
-    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], allow_credentials=True
 )
-DATA_DIR = Path("/app/data")
-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")
@@ -46,103 +40,93 @@ def healthz():
 @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>
-    </form>
   </body>
 </html>"""
-@app.post("/upload")
 async def upload_image(
     file: UploadFile = File(...),
-    top_k: int = Query(5, ge=1, le=20),
 ):
-    content = await file.read()
-    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")

 from __future__ import annotations
 import io
+import os
 from pathlib import Path
 from typing import List
 import gradio as gr
+from fastapi import FastAPI, File, HTTPException, Query, UploadFile
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel, Field
+from PIL import Image
 import tagger as tg
 # -------------------- FastAPI --------------------
+app = FastAPI(
+    title="Image Tagger API",
+    version="1.0.0",
+    description="Generate a caption with BLIP, then return top-K tags derived from that caption.",
 )
+WRITE_SIDECAR = os.getenv("WRITE_SIDECAR", "1") != "0"
+class TagResponse(BaseModel):
+    filename: str = Field(..., examples=["photo.jpg"])
+    caption: str = Field(..., examples=["a lion rests on a rock in the wild"])
+    tags: List[str] = Field(..., examples=[["lion", "rests", "rock", "wild"]])
+@app.on_event("startup")
+def _load_once() -> None:
+    tg.init_models()
 @app.get("/healthz")
 @app.get("/", response_class=HTMLResponse)
 def root():
+    return """
+<!doctype html>
 <html>
+  <head>
+    <meta charset="utf-8" />
+    <title>Image Tagger API</title>
+    <style>
+      body{font-family: system-ui, -apple-system, Segoe UI, Roboto, Ubuntu, sans-serif; max-width: 820px; margin: 48px auto; padding: 0 16px;}
+      .card{border:1px solid #e5e7eb; border-radius:12px; padding:20px;}
+      .btn{background:#111; color:#fff; padding:.6rem 1rem; border-radius:10px; text-decoration:none;}
+      .btn:focus,.btn:hover{opacity:.9}
+      input[type=number]{width:80px;}
+    </style>
+  </head>
+  <body>
+    <h2>🖼️ Image Tagger API</h2>
+    <p>Use <a href="/docs">/docs</a> for Swagger or try the simple UI at <a class="btn" href="/ui">/ui</a>.</p>
+    <div class="card">
+      <h3>Quick upload</h3>
+      <form action="/upload" method="post" enctype="multipart/form-data">
+        <p><input type="file" name="file" accept="image/png,image/jpeg,image/webp" required></p>
+        <p>Top K tags: <input type="number" name="top_k" min="1" max="20" value="5"></p>
+        <p><button class="btn" type="submit">Upload</button></p>
+      </form>
+    </div>
   </body>
 </html>"""
+@app.post("/upload", response_model=TagResponse)
 async def upload_image(
     file: UploadFile = File(...),
+    top_k: int = Query(5, ge=1, le=20, description="How many tags to return"),
 ):
+    try:
+        content = await file.read()
+        img = Image.open(io.BytesIO(content)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid image: {e}")
+    # caption with BLIP
+    caption = tg.caption_image(img)
+    # top-K tags (ensure tagger returns ONLY the list)
+    stem = Path(file.filename).stem
+    tags = tg.caption_to_tags(caption, top_k=top_k)
+    # optional sidecar (same content shape as JSON response)
+    if WRITE_SIDECAR:
+        try:
+            (Path(os.getenv("DATA_DIR", "/app/data"))).mkdir(parents=True, exist_ok=True)
+            (Path(os.getenv("DATA_DIR", "/app/data")) / f"{stem}.json").write_text(
+                TagResponse(filename=file.filename, caption=caption, tags=tags).model_dump_json(indent=2)
+            )
+        except Exception:
+            # ignore filesystem errors; do not fail the request
+            pass
+    return TagResponse(filename=file.filename, caption=caption, tags=tags)
+# -------------------- Gradio (mounted at /ui) --------------------
+def _infer(image: Image.Image, top_k: int):
+    """Wraps the same logic used by the API, but returns simple types
+    so the schema is trivial for Gradio (avoids JSON/dict outputs)."""
+    if image is None:
+        return "", ""
+    cap = tg.caption_image(image)
+    tags = tg.caption_to_tags(cap, top_k=top_k)
+    return cap, ", ".join(tags)
+with gr.Blocks(title="Image Tagger UI") as demo:
+    gr.Markdown("### 🔍 Image → Caption → Tags\nUpload an image → BLIP generates a caption → we extract up to **K** simple tags.")
     with gr.Row():
+        with gr.Column(scale=3):
+            in_img = gr.Image(type="pil", label="Upload image", height=480)
+            k = gr.Slider(1, 20, value=5, step=1, label="Number of tags (K)")
+            submit = gr.Button("Submit", variant="primary")
+            clear = gr.Button("Clear")
+        with gr.Column(scale=2):
+            out_cap = gr.Textbox(label="Generated Caption", lines=2)
+            out_tags = gr.Textbox(label="Tags (comma-separated)", lines=2)
+    submit.click(_infer, inputs=[in_img, k], outputs=[out_cap, out_tags])
+    clear.click(lambda: (None, 5, "", ""), outputs=[in_img, k, out_cap, out_tags])
+# mount Gradio under FastAPI
 app = gr.mount_gradio_app(app, demo, path="/ui")