from fastapi import FastAPI, APIRouter, UploadFile, File, Request
from fastapi.staticfiles import StaticFiles
import shutil
import os
import cv2
import base64
from ultralytics import YOLO
# ─── Load Models 
person_model = YOLO("yolov8n.pt")
statue_model = YOLO("best.pt")
# ─── App Setup 
app = FastAPI()
UPLOAD_FOLDER = "/tmp/uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.mount("/uploads", StaticFiles(directory=UPLOAD_FOLDER), name="uploads")
CONF_THRESHOLD = 0.30
router = APIRouter()
def draw_label(image, label, x1, y1, x2, y2, box_color, text_color):
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.6
    thickness = 2
    (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
    
    if y1 - th - 8 >= 0:
        label_y1 = y1 - th - 8
        label_y2 = y1
        text_y = y1 - 5
    else:
        
        label_y1 = y1
        label_y2 = y1 + th + 8
        text_y = y1 + th + 3
    cv2.rectangle(image, (x1, label_y1), (x1 + tw + 4, label_y2), box_color, -1)
    cv2.putText(image, label, (x1 + 2, text_y), font, font_scale, text_color, thickness)
@app.get("/")
def root():
    return {"message": "AI API is running 🚀"}
@router.post("/predict-image")
async def predict_image(
    request: Request,
    file: UploadFile = File(...)
):
    safe_filename = file.filename.replace(" ", "_")
    file_path = os.path.join(UPLOAD_FOLDER, safe_filename)
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    image = cv2.imread(file_path)
    if image is None:
        return {"error": "Invalid image"}
    detections = []
    person_count = 0
    statue_count = 0
    # ── Person Detection 
    person_results = person_model(file_path)
    for box in person_results[0].boxes:
        cls_id = int(box.cls)
        if cls_id != 0:
            continue
        conf = float(box.conf)
        if conf < CONF_THRESHOLD:
            continue
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        draw_label(image, f"Person {conf:.2f}", x1, y1, x2, y2,
                   box_color=(0, 255, 0), text_color=(0, 0, 0))
        detections.append({
            "type": "person",
            "name": "Person",
            "confidence": round(conf, 4),
            "bbox": [x1, y1, x2, y2]
        })
        person_count += 1
    # ── Statue Detection 
    statue_results = statue_model(file_path)
    for box in statue_results[0].boxes:
        conf = float(box.conf)
        if conf < CONF_THRESHOLD:
            continue
        cls_id = int(box.cls)
        statue_name = statue_results[0].names[cls_id]
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
        draw_label(image, f"{statue_name} {conf:.2f}", x1, y1, x2, y2,
                   box_color=(0, 0, 255), text_color=(255, 255, 255))
        detections.append({
            "type": "statue",
            "name": statue_name,
            "confidence": round(conf, 4),
            "bbox": [x1, y1, x2, y2]
        })
        statue_count += 1
    # ── Save Output Image 
    output_filename = f"output_{safe_filename}"
    output_path = os.path.join(UPLOAD_FOLDER, output_filename)
    cv2.imwrite(output_path, image)
    with open(output_path, "rb") as img_file:
        image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    image_url = f"{request.base_url}uploads/{output_filename}"
    return {
        "total_count": len(detections),
        "persons": person_count,
        "statues": statue_count,
        "output_image_url": image_url,
        "output_image_base64": f"data:image/jpeg;base64,{image_base64}",
        "detections": detections
    }
app.include_router(router)