from fastapi import FastAPI, Request, UploadFile
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
import subprocess
from pathlib import Path

app = FastAPI()

app.mount("/static", FileResponse, name="static")

templates = Jinja2Templates(directory="templates")

class Item(BaseModel):
    file: UploadFile

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

@app.post("/detect", response_class=HTMLResponse)
async def detect_objects(request: Request, item: Item):
    input_file = f"uploads/{item.file.filename}"
    output_file = f"static/output/{item.file.filename}"

    # Save the uploaded file
    with open(input_file, "wb") as file:
        file.write(item.file.file.read())

    # Run YOLOv7 detection and save output
    subprocess.run(["python", "detect.py", "--weights", "model/best.pt", "--img-size", "640", "--source", input_file, "--save-txt", "--save-conf", "--exist-ok", "--output", output_file])

    # Render HTML using Jinja2Templates
    return templates.TemplateResponse(
        "result.html",
        {"request": request, "original_image": f"/static/{item.file.filename}", "output_image": f"/static/output/{item.file.filename}"},
    )


# Start app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)