Update main.py

main.py

@@ -1,24 +1,22 @@
-import os
-import cv2
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
+import torch
+import cv2
 import numpy as np
 import logging
 from io import BytesIO
 import tempfile
-
-from vtoonify_model import Model  # Import VToonify model
+import os
 
 app = FastAPI()
 
-os.makedirs('output', exist_ok=True)
-
-# Initialize VToonify model
+# Load model and necessary components
 model = None
 
 def load_model():
     global model
+    from vtoonify_model import Model
     model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
     model.load_model('cartoon4')
 
@@ -26,7 +24,7 @@ def load_model():
 logging.basicConfig(level=logging.INFO)
 
 @app.post("/upload/")
-async def process_image(file: UploadFile = File(...)):
+async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
     global model
     if model is None:
         load_model()
@@ -44,56 +42,40 @@ async def process_image(file: UploadFile = File(...)):
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Convert BGR to RGB
-    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    # Save the uploaded image temporarily to pass the file path to the model
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
+        cv2.imwrite(temp_file.name, frame_bgr)
+        temp_file_path = temp_file.name
 
-    # Use AnimeGANv3's face detection logic
     try:
-        # Assume AnimeGANv3_src.Convert detects and returns the cropped face
-        det_face = True  # Assume we always want to detect face
-        detected_face, _ = AnimeGANv3_src.Convert(frame_rgb, Style, det_face)
-        
-        if detected_face is None:
-            logging.error("No face detected by AnimeGANv3.")
-            return {"error": "No face detected in the image."}
-
-        # Save the detected face temporarily to pass the file path to the model
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-            cv2.imwrite(temp_file.name, detected_face[:, :, ::-1])  # Convert RGB to BGR for saving
-            temp_file_path = temp_file.name
-
-        try:
-            # Process the detected face using VToonify
-            aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, 0, 0, 0, 0)
-            if aligned_face is None or instyle is None:
-                logging.error("Failed to process the image: No face detected or alignment failed.")
-                return {"error": message}
-
-            processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
-            if processed_image is None:
-                logging.error("Failed to toonify the image.")
-                return {"error": message}
-
-            # Convert the processed image to RGB before returning
-            processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
-
-            # Convert processed image to bytes
-            _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
-
-            # Return the processed image as a streaming response
-            return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
-        
-        finally:
-            # Clean up the temporary file
-            os.remove(temp_file_path)
-
-    except RuntimeError as error:
-        logging.error(f"Error during AnimeGANv3 processing: {error}")
-        return {"error": "Failed to process the image with AnimeGANv3."}
+        # Process the uploaded image using the file path
+        aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, top, bottom, left, right)
+        if aligned_face is None or instyle is None:
+            logging.error("Failed to process the image: No face detected or alignment failed.")
+            return {"error": message}
+
+        processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
+        if processed_image is None:
+            logging.error("Failed to toonify the image.")
+            return {"error": message}
+
+        # Convert the processed image to RGB before returning
+        processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+
+        # Convert processed image to bytes
+        _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
+
+        # Return the processed image as a streaming response
+        return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
+    
+    finally:
+        # Clean up the temporary file
+        os.remove(temp_file_path)
 
+# Mount static files directory
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")
 
+# Define index route
 @app.get("/")
-def index() -> FileResponse:
-    return FileResponse(path="/app/static/index.html", media_type="text/html")
-    
+def index():
+    return FileResponse(path="/app/AB/index.html", media_type="text/html")