Spaces:
Running on Zero
Running on Zero
File size: 3,739 Bytes
47f2ebf | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import base64
import io
import os
from collections import Counter
import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image, ImageDraw, ImageFont
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
MODEL_ID = os.getenv("MODEL_ID", "./model")
CONFIDENCE_THRESHOLD = float(os.getenv("CONFIDENCE_THRESHOLD", "0.35"))
MAX_IMAGE_MB = int(os.getenv("MAX_IMAGE_MB", "15"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
model = RTDetrForObjectDetection.from_pretrained(MODEL_ID)
model.to(DEVICE)
model.eval()
app = FastAPI(
title="Ice Cream Counter API",
version="1.0.0",
description="RT-DETR ice cream product detector and counter. No YOLO.",
)
@app.get("/health")
def health():
return {
"status": "ok",
"model": MODEL_ID,
"device": str(DEVICE),
"threshold": CONFIDENCE_THRESHOLD,
}
def annotate(image, detections):
image = image.copy()
draw = ImageDraw.Draw(image)
try:
font = ImageFont.truetype("DejaVuSans.ttf", 18)
except Exception:
font = ImageFont.load_default()
for d in detections:
x1, y1, x2, y2 = d["box"]
label = f'{d["class"]} {d["confidence"]:.2f}'
draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
bbox = draw.textbbox((x1, y1), label, font=font)
draw.rectangle(bbox, fill="red")
draw.text((x1, y1), label, fill="white", font=font)
return image
@app.post("/predict")
async def predict(
file: UploadFile = File(...),
return_image: bool = True,
):
if not file.content_type or not file.content_type.startswith("image/"):
raise HTTPException(400, "Upload a JPG, PNG, WEBP, or other image file.")
raw = await file.read()
if len(raw) > MAX_IMAGE_MB * 1024 * 1024:
raise HTTPException(
413,
f"Image is too large. Maximum is {MAX_IMAGE_MB} MB.",
)
try:
image = Image.open(io.BytesIO(raw)).convert("RGB")
except Exception as exc:
raise HTTPException(400, f"Could not read image: {exc}")
inputs = processor(images=image, return_tensors="pt")
inputs = {
k: v.to(DEVICE) if torch.is_tensor(v) else v
for k, v in inputs.items()
}
with torch.inference_mode():
outputs = model(**inputs)
target_sizes = torch.tensor(
[[image.height, image.width]],
device=DEVICE,
)
result = processor.post_process_object_detection(
outputs,
threshold=CONFIDENCE_THRESHOLD,
target_sizes=target_sizes,
)[0]
detections = []
counts = Counter()
for score, label, box in zip(
result["scores"],
result["labels"],
result["boxes"],
):
score_value = float(score.item())
label_id = int(label.item())
class_name = model.config.id2label[label_id]
coords = [round(float(x), 2) for x in box.tolist()]
detections.append(
{
"class": class_name,
"confidence": round(score_value, 4),
"box": coords,
}
)
counts[class_name] += 1
response = {
"total": len(detections),
"counts": dict(sorted(counts.items())),
"detections": detections,
}
if return_image:
annotated = annotate(image, detections)
buf = io.BytesIO()
annotated.save(buf, format="JPEG", quality=90)
response["annotated_image_base64"] = base64.b64encode(
buf.getvalue()
).decode("ascii")
return JSONResponse(response)
|