from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from ultralytics import YOLO
import requests
from PIL import Image
import numpy as np
import cv2
import io
import os

app = FastAPI()

# Mount a static directory for serving images
folder = os.path.join(os.getcwd(), "static")
os.makedirs(folder)
app.mount("/static", StaticFiles(directory="static"), name="static")

# Templates for rendering HTML
templates = Jinja2Templates(directory="templates")


def predict_yolo(image_path):
    # Assuming you have a YOLO API endpoint that accepts an image and returns predictions
    with open(image_path, "rb") as file:
        files = {"file": file}
        # Load a model
        model = YOLO('yolov8n.pt')  # pretrained YOLOv8n model

        # Run batched inference on a list of images
        results = model(image_path)  # return a list of Results objects

        # Process results list
        for result in results:
            boxes = result.boxes  # Boxes object for bbox outputs
            # masks = result.masks  # Masks object for segmentation masks outputs
            # keypoints = result.keypoints  # Keypoints object for pose outputs
            # probs = result.probs  # Probs object for classification outputs
        predictions = boxes.json()
    return predictions


def draw_boxes(image, boxes):
    for box in boxes:
        x, y, w, h = box["bbox"]
        cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 0), 2)
    return image


@app.post("/uploadfile/")
async def create_upload_file(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    
    # Save the image to a static directory
    save_path = f"static/{file.filename}"
    image.save(save_path)
    
    # Perform YOLO prediction
    predictions = predict_yolo(save_path)
    
    # Draw bounding boxes on the image
    image_np = np.array(image)
    image_with_boxes = draw_boxes(image_np, predictions)
    
    # Save the image with bounding boxes
    image_with_boxes_path = f"static/{file.filename.split('.')[0]}_with_boxes.jpg"
    cv2.imwrite(image_with_boxes_path, cv2.cvtColor(image_with_boxes, cv2.COLOR_RGB2BGR))
    
    # Render the HTML with the image and bounding boxes
    return templates.TemplateResponse("prediction.html", {"request": file, "image_path": image_with_boxes_path})


@app.get("/")
async def read_root():
    return {"message": "Hello, this is a YOLO prediction API using FastAPI!"}