Image_Tagger / app.py
stephenebert's picture
Update app.py
ae57f86 verified
raw
history blame
3.33 kB
from __future__ import annotations
# FastAPI app exposing:
# GET /healthz -> {"ok": true}
# GET / -> tiny upload form (links to /docs)
# POST /upload -> {"filename","caption","tags"}
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
from typing import List
from pathlib import Path
from PIL import Image
import io, json
from tagger import tag_pil_image, CAP_TAG_DIR # tagger writes ./data/<stem>.json
app = FastAPI(title="Image Tagger API", version="0.3.0")
class TagOut(BaseModel):
filename: str
caption: str
tags: List[str]
@app.get("/", response_class=HTMLResponse)
def home() -> str:
# Simple HTML so you can upload from the root page
return """
<html><head><meta charset="utf-8"><title>Image Tagger API</title></head>
<body style="font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,sans-serif;max-width:720px;margin:40px auto;padding:0 16px">
<h2>Image Tagger API</h2>
<p>Use <a href="/docs">/docs</a> for full Swagger UI, or upload here:</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" style="padding:.6rem 1rem;border-radius:10px;border:1px solid #ddd;background:#111;color:#fff">Upload</button>
</form>
</body></html>
"""
@app.get("/healthz")
def healthz():
return {"ok": True}
@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"),
# kept for backward compatibility with earlier UI; tagger ignores them
nouns: bool = Query(True, description="(ignored)"),
adjs: bool = Query(True, description="(ignored)"),
verbs: bool = Query(True, description="(ignored)"),
):
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")
# Load image
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")
# Create a sensible stem for sidecar filename
stem = Path(file.filename).stem or "upload"
# Generate tags (and sidecar JSON with caption written by tagger.py)
try:
tags = tag_pil_image(
img,
stem,
top_k=top_k,
keep_nouns=nouns,
keep_adjs=adjs,
keep_verbs=verbs,
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Tagging failed: {e}")
# Read caption from the sidecar if present
caption = ""
sidecar = CAP_TAG_DIR / f"{stem}.json"
if sidecar.exists():
try:
caption = json.loads(sidecar.read_text()).get("caption", "")
except Exception:
caption = ""
return JSONResponse({"filename": file.filename, "caption": caption, "tags": tags})