Spaces:

stephenebert
/

Image_Tagger

Sleeping

App Files Files Community

stephenebert commited on Aug 13, 2025

Commit

a6de95a

verified ·

1 Parent(s): 0b1423b

Update app.py

Browse files

Files changed (1) hide show

app.py +26 -45

app.py CHANGED Viewed

@@ -7,17 +7,18 @@ 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
 )
@@ -26,33 +27,30 @@ 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>
@@ -61,61 +59,44 @@ def root():
   </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")

 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
+# 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(
     CORSMiddleware,
     allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], allow_credentials=True
 )
 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>
   </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
+    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")