pbmarcy/acne-detector-api

1from fastapi import FastAPI, File, UploadFile2from fastapi.responses import JSONResponse3from ultralytics import YOLO4import torch5 6from PIL import Image7import numpy as np8import tensorflow as tf9import io10 11app = FastAPI()12 13# ✅ Load YOLO model (trust source - force full load)14model = YOLO("best.pt")15 16# ✅ Load TFLite severity model17severity_interpreter = tf.lite.Interpreter(model_path="severity_model.tflite")18severity_interpreter.allocate_tensors()19input_details = severity_interpreter.get_input_details()20output_details = severity_interpreter.get_output_details()21 22SEVERITY_LABELS = ["Mild", "Moderate", "Severe"]23 24@app.get("/")25def root():26    return {"message": "Acne Detector API is running"}27 28@app.post("/detect")29async def detect(file: UploadFile = File(...)):30    try:31        image_bytes = await file.read()32        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")33 34        results = model.predict(image)[0]35        boxes = results.boxes.xyxy.tolist()36        confs = results.boxes.conf.tolist()37        classes = results.boxes.cls.tolist()38 39        detections = []40        for box, conf, cls in zip(boxes, confs, classes):41            x_min, y_min, x_max, y_max = map(int, box)42            cropped = image.crop((x_min, y_min, x_max, y_max)).resize((224, 224))43            input_data = np.expand_dims(np.array(cropped) / 255.0, axis=0).astype(np.float32)44 45            severity_interpreter.set_tensor(input_details[0]['index'], input_data)46            severity_interpreter.invoke()47            output_data = severity_interpreter.get_tensor(output_details[0]['index'])48            severity = SEVERITY_LABELS[np.argmax(output_data)]49 50            detections.append({51                "box": box,52                "confidence": round(conf, 3),53                "class_id": int(cls),54                "severity": severity55            })56 57        return JSONResponse(content={58            "detections": detections,59            "count": len(detections)60        })61 62    except Exception as e:63        return JSONResponse(content={"error": str(e)}, status_code=500)64