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("yolo_modeln11_1502.pt")

class ImageRequest(BaseModel):
    image: str

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

@app.post("/predict")
async def predict(request: ImageRequest):
    # Check if the image string is empty
    if not request.image:
        return {"error": "Empty image string"}
    
    try:
        # Attempt to decode Base64 string to bytes
        image_bytes = base64.b64decode(request.image, validate=True)
    except Exception as e:
        return {"error": "Invalid Base64 string", "details": str(e)}
    
    # Ensure the decoded bytes are not empty
    if not image_bytes:
        return {"error": "Decoded data is empty"}
    
    np_arr = np.frombuffer(image_bytes, np.uint8)
    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
    if image is None:
        return {"error": "Invalid image"}
    
    
    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())

        # Counting occurrences of each defect
        if class_name in class_counters:
            class_counters[class_name] += 1
        else:
            class_counters[class_name] = 1
        
        key = f"{class_name}{class_counters[class_name]}" if class_counters[class_name] > 1 else class_name
        json_result[key] = [x1, y1, x2, y2]
    
    # Using model's summary output if available
    if hasattr(result, "summary") and isinstance(result.summary, dict):
        statistics_summary = result.summary
    else:
        statistics_summary = {name: count for name, count in class_counters.items()}
    
    return {"statistics": statistics_summary, "detections": json_result}

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