KineticHonda/yolo_API

1import base642import cv23import numpy as np4import threading5import traceback6 7from fastapi import FastAPI, Request, HTTPException8from fastapi.responses import HTMLResponse, JSONResponse9from fastapi.middleware.cors import CORSMiddleware10 11# Import YOLO from ultralytics (YOLO11)12from ultralytics import YOLO13 14app = FastAPI()15 16# Add CORS middleware17app.add_middleware(18    CORSMiddleware,19    allow_origins=["*", "https://daas-mobile-application.vercel.app", "http://localhost:4200"],20    allow_credentials=True,21    allow_methods=["*"],22    allow_headers=["*"],23)24 25# Path to your custom YOLO11 model file (best.pt) - CHANGED FOR HUGGING FACE26model_path = "best.pt"  # Direct path since file will be in root directory27 28# Load the YOLO11 model using the ultralytics API on CPU29model = YOLO(model_path).to("cpu") 30 31# Global lock for model inference32model_lock = threading.Lock()33 34# Define class ID to name mapping35CLASS_NAMES = {36    0: "Broken Hinge",37    1: "Crack",38    2: "Damaged Keys",39    3: "Damaged Screen",40    4: "Display Issues",41    5: "Laptop",42    6: "Scratch"43}44 45@app.get("/")46async def root():47    """Health check endpoint."""48    return {"message": "DAAS YOLO API is running!", "status": "healthy"}49 50@app.post("/predict")51async def predict(request: Request):52    try:53        data = await request.json()54        image_data = data.get("image")55        if not image_data:56            raise HTTPException(status_code=400, detail="No image data provided")57 58        if "," in image_data:59            header, encoded = image_data.split(",", 1)60        else:61            encoded = image_data62 63        img_bytes = base64.b64decode(encoded)64        np_arr = np.frombuffer(img_bytes, np.uint8)65        image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)66        if image is None:67            raise HTTPException(status_code=400, detail="Could not decode image")68 69        with model_lock:70            results = model(image, conf=0.10)71            annotated_image = results[0].plot()72 73        detected_classes = []74        boxes = getattr(results[0], "boxes", None)75        if boxes and len(boxes) > 0:76            for box in boxes:77                class_id = int(box.cls.item())  # Convert tensor to integer78                confidence = float(box.conf.item())  # Convert tensor to float79                class_name = CLASS_NAMES.get(class_id, "Unknown")  # Get class name from mapping80                detected_classes.append({"class_id": class_id, "class_name": class_name, "confidence": confidence})81        else:82            cv2.putText(83                annotated_image,84                "No damages detected",85                (10, 30),86                cv2.FONT_HERSHEY_SIMPLEX,87                1,88                (0, 0, 255),89                2,90                cv2.LINE_AA91            )92 93        success, encoded_image = cv2.imencode(".jpg", annotated_image)94        if not success:95            raise HTTPException(status_code=500, detail="Could not encode image")96        annotated_base64 = base64.b64encode(encoded_image).decode("utf-8")97        data_url = "data:image/jpeg;base64," + annotated_base6498 99        return JSONResponse(content={"image": data_url, "detected_classes": detected_classes})100 101    except Exception as e:102        print("Exception in /predict:\n", traceback.format_exc())103        raise HTTPException(status_code=500, detail=str(e))104    105if __name__ == "__main__":106    import uvicorn107    # Run the app on port 7860 (required for Hugging Face Spaces)108    uvicorn.run(app, host="0.0.0.0", port=7860)