Spaces:
Running
Running
File size: 1,329 Bytes
a9db693 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | from __future__ import annotations
from fastapi import APIRouter, File, HTTPException, Response, UploadFile
from ..ml.segment import segment_image
from ..utils.images import load_pil
router = APIRouter()
@router.post("")
async def segment_ep(file: UploadFile = File(...), output: str = "json"):
try:
data = await file.read()
img = load_pil(data)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc))
res = segment_image(img)
if output == "overlay":
return Response(content=res.overlay_png, media_type="image/png")
if output == "mask":
return Response(content=res.mask_png, media_type="image/png")
return {
"attributes": res.attributes,
"shape": list(res.mask.shape),
"categories": [
"wall", "window", "door", "roof", "balcony",
"column", "sky", "vegetation", "ground", "other",
],
}
@router.post("/color")
async def palette_ep(file: UploadFile = File(...)):
from ..ml.color import palette
try:
data = await file.read()
img = load_pil(data)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc))
items = palette(img, k=5)
return [{"hex": e.rgb_hex, "rgb": list(e.rgb), "share": e.share} for e in items]
|