Update main.py

main.py

@@ -1,23 +1,24 @@
-from fastapi import FastAPI, File, UploadFile
+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
-import os
+import AnimeGANv3_src  # Assuming this module contains the face detection logic
+from vtoonify_model import Model  # Import VToonify model
 
 app = FastAPI()
 
-# Load model and necessary components
+os.makedirs('output', exist_ok=True)
+
+# Initialize VToonify model
 model = None
-face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
 
 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')
 
@@ -25,7 +26,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(...), Style: str = Form(...)):
     global model
     if model is None:
         load_model()
@@ -43,52 +44,55 @@ async def process_image(file: UploadFile = File(...)):
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Detect faces in the image
-    gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
-    if len(faces) == 0:
-        logging.error("No faces detected in the image.")
-        return {"error": "No faces detected in the image."}
-
-    # Use the first detected face
-    (x, y, w, h) = faces[0]
-    top, bottom, left, right = y, y + h, x, x + w
-
-    # 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
+    # Convert BGR to RGB
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
 
+    # Use AnimeGANv3's face detection logic
     try:
-        # 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")
+        # 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."}
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
-# Define index route
 @app.get("/")
-def index():
-    return FileResponse(path="/app/AB/index.html", media_type="text/html")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")