Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -1,53 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
|
| 2 |
-
from fastapi.responses import JSONResponse
|
| 3 |
from pydantic import BaseModel
|
| 4 |
from typing import List
|
| 5 |
from pathlib import Path
|
| 6 |
from PIL import Image
|
| 7 |
import io, json
|
| 8 |
|
| 9 |
-
from tagger import tag_pil_image, CAP_TAG_DIR
|
| 10 |
|
| 11 |
app = FastAPI(title="Image Tagger API", version="0.3.0")
|
| 12 |
|
|
|
|
| 13 |
class TagOut(BaseModel):
|
| 14 |
filename: str
|
| 15 |
caption: str
|
| 16 |
tags: List[str]
|
| 17 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
@app.get("/healthz")
|
| 19 |
def healthz():
|
| 20 |
return {"ok": True}
|
| 21 |
|
|
|
|
| 22 |
@app.post("/upload", response_model=TagOut)
|
| 23 |
async def upload(
|
| 24 |
-
file: UploadFile = File(
|
| 25 |
-
top_k: int = Query(5, ge=1, le=20, description="
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
|
|
|
| 29 |
):
|
| 30 |
-
if file.content_type not in {"image/png", "image/jpeg"}:
|
| 31 |
-
raise HTTPException(415, "Only PNG or
|
| 32 |
|
|
|
|
| 33 |
try:
|
| 34 |
data = await file.read()
|
| 35 |
img = Image.open(io.BytesIO(data)).convert("RGB")
|
| 36 |
except Exception:
|
| 37 |
-
raise HTTPException(400, "Could not decode image")
|
| 38 |
|
|
|
|
| 39 |
stem = Path(file.filename).stem or "upload"
|
| 40 |
-
tags = tag_pil_image(
|
| 41 |
-
img, stem,
|
| 42 |
-
top_k=top_k, keep_nouns=nouns, keep_adjs=adjs, keep_verbs=verbs
|
| 43 |
-
)
|
| 44 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 45 |
caption = ""
|
| 46 |
-
|
| 47 |
-
if
|
| 48 |
try:
|
| 49 |
-
caption = json.loads(
|
| 50 |
except Exception:
|
| 51 |
-
|
| 52 |
|
| 53 |
return JSONResponse({"filename": file.filename, "caption": caption, "tags": tags})
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
# FastAPI app exposing:
|
| 4 |
+
# GET /healthz -> {"ok": true}
|
| 5 |
+
# GET / -> tiny upload form (links to /docs)
|
| 6 |
+
# POST /upload -> {"filename","caption","tags"}
|
| 7 |
+
|
| 8 |
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
|
| 9 |
+
from fastapi.responses import HTMLResponse, JSONResponse
|
| 10 |
from pydantic import BaseModel
|
| 11 |
from typing import List
|
| 12 |
from pathlib import Path
|
| 13 |
from PIL import Image
|
| 14 |
import io, json
|
| 15 |
|
| 16 |
+
from tagger import tag_pil_image, CAP_TAG_DIR # tagger writes ./data/<stem>.json
|
| 17 |
|
| 18 |
app = FastAPI(title="Image Tagger API", version="0.3.0")
|
| 19 |
|
| 20 |
+
|
| 21 |
class TagOut(BaseModel):
|
| 22 |
filename: str
|
| 23 |
caption: str
|
| 24 |
tags: List[str]
|
| 25 |
|
| 26 |
+
|
| 27 |
+
@app.get("/", response_class=HTMLResponse)
|
| 28 |
+
def home() -> str:
|
| 29 |
+
# Simple HTML so you can upload from the root page
|
| 30 |
+
return """
|
| 31 |
+
<html><head><meta charset="utf-8"><title>Image Tagger API</title></head>
|
| 32 |
+
<body style="font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,sans-serif;max-width:720px;margin:40px auto;padding:0 16px">
|
| 33 |
+
<h2>Image Tagger API</h2>
|
| 34 |
+
<p>Use <a href="/docs">/docs</a> for full Swagger UI, or upload here:</p>
|
| 35 |
+
<form action="/upload" method="post" enctype="multipart/form-data" style="display:grid;gap:12px">
|
| 36 |
+
<input type="file" name="file" accept="image/png,image/jpeg,image/webp" required />
|
| 37 |
+
<label>top_k:
|
| 38 |
+
<input type="number" name="top_k" value="5" min="1" max="20">
|
| 39 |
+
</label>
|
| 40 |
+
<button type="submit" style="padding:.6rem 1rem;border-radius:10px;border:1px solid #ddd;background:#111;color:#fff">Upload</button>
|
| 41 |
+
</form>
|
| 42 |
+
</body></html>
|
| 43 |
+
"""
|
| 44 |
+
|
| 45 |
+
|
| 46 |
@app.get("/healthz")
|
| 47 |
def healthz():
|
| 48 |
return {"ok": True}
|
| 49 |
|
| 50 |
+
|
| 51 |
@app.post("/upload", response_model=TagOut)
|
| 52 |
async def upload(
|
| 53 |
+
file: UploadFile = File(...),
|
| 54 |
+
top_k: int = Query(5, ge=1, le=20, description="Max number of tags"),
|
| 55 |
+
# kept for backward compatibility with earlier UI; tagger ignores them
|
| 56 |
+
nouns: bool = Query(True, description="(ignored)"),
|
| 57 |
+
adjs: bool = Query(True, description="(ignored)"),
|
| 58 |
+
verbs: bool = Query(True, description="(ignored)"),
|
| 59 |
):
|
| 60 |
+
if file.content_type not in {"image/png", "image/jpeg", "image/webp"}:
|
| 61 |
+
raise HTTPException(status_code=415, detail="Only PNG, JPEG, or WebP images are supported")
|
| 62 |
|
| 63 |
+
# Load image
|
| 64 |
try:
|
| 65 |
data = await file.read()
|
| 66 |
img = Image.open(io.BytesIO(data)).convert("RGB")
|
| 67 |
except Exception:
|
| 68 |
+
raise HTTPException(status_code=400, detail="Could not decode image")
|
| 69 |
|
| 70 |
+
# Create a sensible stem for sidecar filename
|
| 71 |
stem = Path(file.filename).stem or "upload"
|
|
|
|
|
|
|
|
|
|
|
|
|
| 72 |
|
| 73 |
+
# Generate tags (and sidecar JSON with caption written by tagger.py)
|
| 74 |
+
try:
|
| 75 |
+
tags = tag_pil_image(
|
| 76 |
+
img,
|
| 77 |
+
stem,
|
| 78 |
+
top_k=top_k,
|
| 79 |
+
keep_nouns=nouns,
|
| 80 |
+
keep_adjs=adjs,
|
| 81 |
+
keep_verbs=verbs,
|
| 82 |
+
)
|
| 83 |
+
except Exception as e:
|
| 84 |
+
raise HTTPException(status_code=500, detail=f"Tagging failed: {e}")
|
| 85 |
+
|
| 86 |
+
# Read caption from the sidecar if present
|
| 87 |
caption = ""
|
| 88 |
+
sidecar = CAP_TAG_DIR / f"{stem}.json"
|
| 89 |
+
if sidecar.exists():
|
| 90 |
try:
|
| 91 |
+
caption = json.loads(sidecar.read_text()).get("caption", "")
|
| 92 |
except Exception:
|
| 93 |
+
caption = ""
|
| 94 |
|
| 95 |
return JSONResponse({"filename": file.filename, "caption": caption, "tags": tags})
|