Update main.py

main.py

@@ -1,22 +1,23 @@
-from fastapi import FastAPI, File, UploadFile, Form
+import os
+import cv2
+from fastapi import FastAPI, File, UploadFile
 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
+from mtcnn import MTCNN
+from vtoonify_model import Model  # Import VToonify model
 
 app = FastAPI()
 
-# Load model and necessary components
+# Initialize the VToonify model and MTCNN detector
 model = None
+detector = MTCNN()
 
 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')
 
@@ -24,7 +25,7 @@ def load_model():
 logging.basicConfig(level=logging.INFO)
 
 @app.post("/upload/")
-async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
+async def process_image(file: UploadFile = File(...)):
     global model
     if model is None:
         load_model()
@@ -32,24 +33,38 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
     # Read the uploaded image file
     contents = await file.read()
 
-    # Convert the uploaded image to numpy array
+    # Convert the uploaded image to a numpy array
     nparr = np.frombuffer(contents, np.uint8)
     frame_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)  # Read as BGR format by default
-    
+
     if frame_bgr is None:
         logging.error("Failed to decode the image.")
         return {"error": "Failed to decode the image. Please ensure the file is a valid image format."}
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Save the uploaded image temporarily to pass the file path to the model
+    # Convert BGR to RGB for MTCNN
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+
+    # Detect faces using MTCNN
+    results = detector.detect_faces(frame_rgb)
+
+    if len(results) == 0:
+        logging.error("No faces detected in the image.")
+        return {"error": "No faces detected in the image."}
+
+    # Use the first detected face
+    x, y, width, height = results[0]['box']
+    cropped_face = frame_rgb[y:y+height, x:x+width]
+
+    # Save the cropped 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, frame_bgr)
+        cv2.imwrite(temp_file.name, cv2.cvtColor(cropped_face, cv2.COLOR_RGB2BGR))
         temp_file_path = temp_file.name
 
     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)
+        # Process the cropped 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}
@@ -67,15 +82,15 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
 
         # 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")
+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")