Update main.py

main.py

@@ -7,56 +7,41 @@ import numpy as np
 import cv2
 import torch
 from vtoonify_model import Model
+import shutil
+import os
 
 app = FastAPI()
-
-# Load the model
 model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
-exstyle, message = model.load_model("cartoon1")
-
-class ImageRequest(BaseModel):
-    image_file: UploadFile = File(...)
-
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
-import numpy as np
-import cv2
-import torch
-from vtoonify_model import Model
 
-app = FastAPI()
-
-# Load the model
-model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
-exstyle, message = model.load_model("cartoon1")
-
-@app.post("/upload/")
-async def toonify_image(image_file: UploadFile = File(...)):
+@app.post("/upload")
+async def upload_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
     try:
-        # Read the contents of the uploaded file
-        contents = await image_file.read()
-        # Determine the file extension from the content type
-        ext = image_file.content_type.split("/")[-1]
-        # Construct a unique filename with the extension
-        filename = f"uploaded_image.{ext}"
+        with open("temp.jpg", "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
         
-        # Convert the file contents to a numpy array
-        nparr = np.frombuffer(contents, np.uint8)
-        # Decode the image using OpenCV
-        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        # Perform image processing using the model
+        result_image, _, _ = model.detect_and_align_image("temp.jpg", top, bottom, left, right)
         
-        # Perform image processing with the model
-        aligned_face, instyle, message = model.detect_and_align_image(img, 200, 200, 200, 200)  # Hardcoded values
-        toonified_img, message = model.image_toonify(aligned_face, instyle, exstyle, style_degree=0.5, style_type="cartoon1")
+        # Save the result image temporarily
+        result_path = "result.jpg"
+        cv2.imwrite(result_path, result_image)
         
-        # Return the processed image and filename
-        return {"toonified_image": toonified_img, "message": message, "filename": filename}
+        # 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 an error response if an exception occurs
-        return JSONResponse(status_code=500, content={"error": str(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")
+