Spaces:

stephenebert
/

Image_Tagger

Sleeping

App Files Files Community

stephenebert commited on Aug 13, 2025

Commit

3452c0c

verified ·

1 Parent(s): 7b6684e

Update app.py

Browse files

Files changed (1) hide show

app.py +104 -68

app.py CHANGED Viewed

@@ -1,86 +1,122 @@
 from __future__ import annotations
-# FastAPI REST API + Gradio UI at /
-# Endpoints:
-#   GET  /healthz
-#   POST /upload   -> {filename, caption, tags}
-# UI:
-#   /  (upload image, choose top_k, see caption + tags)
-# Docs:
-#   /docs
-from fastapi import FastAPI, File, HTTPException, Query, UploadFile
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-from typing import List
-from pathlib import Path
-from PIL import Image
 import io
 import gradio as gr
-from tagger import tag_pil_image  # returns (caption: str, tags: List[str])
-app = FastAPI(title="Image Tagger API", version="0.4.3")
-# ---------- Pydantic model ----------
-class TagOut(BaseModel):
-    filename: str
-    caption: str
-    tags: List[str]
-# ---------- Health ----------
-@app.get("/healthz")
 def healthz():
     return {"ok": True}
-# ---------- REST endpoint ----------
-@app.post("/upload", response_model=TagOut)
-async def upload(
     file: UploadFile = File(...),
-    top_k: int = Query(5, ge=1, le=20, description="Max number of tags"),
 ):
-    if file.content_type not in {"image/png", "image/jpeg", "image/webp"}:
-        raise HTTPException(
-            status_code=415, detail="Only PNG, JPEG, or WebP images are supported"
-        )
-    try:
-        data = await file.read()
-        img = Image.open(io.BytesIO(data)).convert("RGB")
-    except Exception:
-        raise HTTPException(status_code=400, detail="Could not decode image")
-    stem = Path(file.filename).stem or "upload"
-    try:
-        caption, tags = tag_pil_image(img, stem, top_k=top_k)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Tagging failed: {e}")
-    return JSONResponse({"filename": file.filename, "caption": caption, "tags": tags})
-# ---------- Gradio UI at root ----------
-def _ui_tag(image: Image.Image, top_k: int):
-    if image is None:
-        return "", ""
-    caption, tags = tag_pil_image(image.convert("RGB"), "upload", top_k=top_k)
-    return caption, ", ".join(tags)
-demo = gr.Interface(
-    fn=_ui_tag,
-    inputs=[
-        gr.Image(type="pil", label="Upload image"),
-        gr.Slider(1, 20, value=5, step=1, label="Top-k tags"),
-    ],
-    outputs=[
-        gr.Textbox(label="Caption", lines=2),
-        gr.Textbox(label="Tags (comma-separated)", lines=2),
-    ],
-    flagging_mode="never",
-    title="Image Tagger",
-    description="Upload an image to get a caption and top-k tags. Programmatic API at /docs.",
-)
-# Mount Gradio on the same FastAPI app at root (/) to avoid redirects
-app = gr.mount_gradio_app(app, demo, path="/")

+# 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.responses import HTMLResponse, JSONResponse, RedirectResponse
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
 import gradio as gr
+# your tagger module (already loads BLIP and returns top-k tags + writes sidecar)
+import tagger as tg
+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)
+# ---------- Helpers ----------
+def _caption_with_tagger(img: Image.Image) -> str:
+    """Use the BLIP model objects loaded by tagger.py to generate a caption."""
+    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:
+        # Caption is optional; only tags are required per your latest request.
+        return ""
+# ---------- FastAPI endpoints ----------
+@APP.get("/healthz")
 def healthz():
     return {"ok": True}
+@APP.get("/", response_class=HTMLResponse)
+def root():
+    # Keep this simple (no schema generation). Link to /docs and /ui.
+    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>Upload via <a href="/docs">/docs</a> or try 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>
+    </form>
+  </body>
+</html>"""
+@APP.post("/upload")
+async def upload_image(
     file: UploadFile = File(...),
+    top_k: int = Query(5, ge=1, le=20),
 ):
+    # Read file into PIL
+    content = await file.read()
+    img = Image.open(io.BytesIO(content)).convert("RGB")
+    stem = Path(file.filename).stem
+    # Get tags (tagger will also write a JSON sidecar under tg.CAP_TAG_DIR)
+    tags: List[str] = tg.tag_pil_image(img, stem, top_k=top_k)
+    # Optional caption (doesn't affect tags)
+    caption = _caption_with_tagger(img)
+    payload = {
+        "filename": file.filename,
+        "caption": caption,
+        "tags": tags,
+    }
+    # Also store a copy under /app/data 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 (mounted 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)
+    payload = {"filename": "ui_upload", "caption": caption, "tags": tags}
+    return caption, ", ".join(tags), json.dumps(payload, indent=2)
+with gr.Blocks(css="""
+  .gradio-container {max-width: 860px !important}
+""") as demo:
+    gr.Markdown("## Image Tagger — UI\nUpload an image, choose `top_k`, and get tags.")
+    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")
+    with gr.Row():
+        out_caption = gr.Textbox(label="Caption", lines=2)
+    with gr.Row():
+        out_tags = gr.Textbox(label="Tags (comma-separated)", lines=2)
+    out_json = gr.Textbox(label="Raw JSON", lines=10)
+    run.click(_gr_predict, inputs=[in_img, k], outputs=[out_caption, out_tags, out_json])
+APP = gr.mount_gradio_app(APP, demo, path="/ui")