from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
import cv2
import shutil
from model import Model  # Import your model class here

app = FastAPI()
model = Model(device='cpu')  # Initialize your model with the appropriate device

@app.post("/upload/")
async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
    # Save the uploaded image locally
    with open("uploaded_image.jpg", "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # Load the model (assuming 'cartoon1' is always used)
    exstyle, load_info = model.load_model('cartoon1')

    # Process the uploaded image
    # Replace these values with the actual bounding box coordinates
    aligned_face, instyle, message = model.detect_and_align_image("uploaded_image.jpg", top, bottom, left, right)
    
    # Process the aligned face further if needed
    # For example, you can pass it to the image_toonify method
    style_degree = 0.5
    style_type = 'cartoon1'  # Adjust this based on the actual style type
    result_image, message = model.image_toonify(aligned_face, instyle, exstyle, style_degree, style_type)

    # Save the result image locally
    result_image_path = "result_image.jpg"
    cv2.imwrite(result_image_path, result_image)

    # Return the processed image
    return FileResponse(result_image_path, media_type="image/jpeg", headers={"Content-Disposition": "attachment; filename=result_image.jpg"})