from fastapi import FastAPI
import uvicorn
import base64
import cv2
import numpy as np
from ultralytics import YOLO
from datetime import datetime
from pydantic import BaseModel

app = FastAPI()

model = YOLO("pcb_component_detection_best.pt")

class ImageRequest(BaseModel):
    image: str

@app.get("/")
async def root():
    current_time = datetime.now().isoformat()
    return {"message": "PCB components API works", "time": current_time}

@app.post("/predict")
async def predict(request: ImageRequest):
    # Decode Base64
    image_bytes = base64.b64decode(request.image)
    np_arr = np.frombuffer(image_bytes, np.uint8)
    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
    if image is None:
        return {"error": "Invalid image"}
    
    # Inference
    results = model.predict(image)
    result = results[0]

    # Response
    json_result = {}
    class_counters = {}
    
    for box in result.boxes:
        class_id = int(box.cls[0])
        class_name = result.names[class_id]
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())

        if class_name in class_counters:
            class_counters[class_name] += 1
            key = f"{class_name}_{class_counters[class_name]}"
        else:
            class_counters[class_name] = 1
            key = class_name
        
        json_result[key] = [x1, y1, x2, y2]
    
    return json_result

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)