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

app = FastAPI()

# Initialize templates
templates = Jinja2Templates(directory="templates")

# Load YOLO model
current_directory = os.getcwd()
model_path = current_directory + '\\datasets\\runs\\detect\\train8\\weights\\last.pt'
model = YOLO(model_path)

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/upload")
async def upload_image(file: UploadFile = File(...)):
    # Read the uploaded file
    contents = await file.read()
    
    # Convert to numpy array
    nparr = np.frombuffer(contents, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # Run YOLO prediction
    results = model(img)
    
    # Draw bounding boxes
    annotated_img = results[0].plot()
    
    # Convert to base64
    _, img_encoded = cv2.imencode('.jpg', annotated_img)
    img_base64 = base64.b64encode(img_encoded).decode('utf-8')
    
    return {"image": img_base64}

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)