from fastapi import FastAPI, File, UploadFile,Form
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
import numpy as np
import cv2
import torch
from vtoonify_model import Model
import shutil
import os

app = FastAPI()
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')

@app.post("/upload/")
async def upload_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
    try:
        with open("temp.jpg", "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Perform image processing using the model
        result_image, _, _ = model.detect_and_align_image("temp.jpg", top, bottom, left, right)
        
        # Save the result image temporarily
        result_path = "result.jpg"
        cv2.imwrite(result_path, result_image)
        
        # Return the result image
        with open(result_path, "rb") as result_buffer:
            result_image_bytes = result_buffer.read()
        
        # Remove temporary files
        os.remove("temp.jpg")
        os.remove(result_path)
        
        return {"result_image": result_image_bytes}
    
    except Exception as e:
        return {"error": str(e)}

app.mount("/", StaticFiles(directory="AB", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/AB/index.html", media_type="text/html")