VT3 / main.py
Ashrafb's picture
Update main.py
3358ca8 verified
raw
history blame
1.55 kB
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"})