Update main.py

main.py

@@ -1,32 +1,30 @@
-import os
-import cv2
-from fastapi import FastAPI, File, UploadFile
+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 mtcnn import MTCNN
-from vtoonify_model import Model  # Import VToonify model
-import torch
+import os
 
 app = FastAPI()
 
-# Initialize logging
-logging.basicConfig(level=logging.INFO)
-
-# Initialize the VToonify model and MTCNN detector
+# Load model and necessary components
 model = None
-detector = MTCNN(min_face_size=20, scale_factor=0.709)
 
 def load_model():
     global model
+    from vtoonify_model import Model
     model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
-    model.load_model('cartoon1')
+    model.load_model('cartoon4')
+
+# Configure logging
+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()
@@ -34,65 +32,50 @@ async def process_image(file: UploadFile = File(...)):
     # Read the uploaded image file
     contents = await file.read()
 
-    # Convert the uploaded image to a numpy array
+    # Convert the uploaded image to 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}")
 
-    # 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)
-    logging.info(f"Detection results: {results}")
-
-    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=True, suffix=".jpg") as temp_file:
-        cv2.imwrite(temp_file.name, cv2.cvtColor(cropped_face, cv2.COLOR_RGB2BGR))
+    # 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
 
-        try:
-            # 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}
-
-            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}
+    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}
 
-            # Convert the processed image to RGB before returning
-            processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+        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 processed image to bytes
-            _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
+        # Convert the processed image to RGB before returning
+        processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
 
-            # Return the processed image as a streaming response
-            return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
+        # Convert processed image to bytes
+        _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
 
-        except Exception as e:
-            logging.error(f"Error during processing: {e}")
-            return {"error": str(e)}
+        # 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="AB")
+app.mount("/", StaticFiles(directory="AB", html=True), name="static")
 
 # Define index route
 @app.get("/")
-def index() -> FileResponse:
+def index():
     return FileResponse(path="/app/AB/index.html", media_type="text/html")