Update main.py

main.py

@@ -8,6 +8,7 @@ import logging
 from io import BytesIO
 import tempfile
 import os
+from insightface.app import FaceAnalysis
 
 app = FastAPI()
 
@@ -20,9 +21,53 @@ def load_model():
     model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
     model.load_model('cartoon4')
 
+# Initialize the InsightFace model for face detection
+face_detector = FaceAnalysis(allowed_modules=['detection'])
+face_detector.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(640, 640))
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
+def detect_and_crop_face(image, padding=0.6):
+    # Get original dimensions
+    orig_h, orig_w = image.shape[:2]
+    
+    # Resize the image for detection
+    resized_image = cv2.resize(image, (640, 640))
+    
+    # Detect faces on the resized image
+    faces = face_detector.get(resized_image)
+    
+    # If faces are detected, sort by x-coordinate and select the leftmost face
+    if faces:
+        faces = sorted(faces, key=lambda face: face.bbox[0])
+        face = faces[0]  # Select the leftmost face
+        bbox = face.bbox.astype(int)
+        
+        # Calculate scaling factors
+        h_scale = orig_h / 640
+        w_scale = orig_w / 640
+        
+        # Map the bounding box to the original image size
+        x1, y1, x2, y2 = bbox
+        x1 = int(x1 * w_scale)
+        y1 = int(y1 * h_scale)
+        x2 = int(x2 * w_scale)
+        y2 = int(y2 * h_scale)
+        
+        # Calculate padding
+        width = x2 - x1
+        height = y2 - y1
+        x1 = max(0, x1 - int(padding * width))
+        y1 = max(0, y1 - int(padding * height))
+        x2 = min(orig_w, x2 + int(padding * width))
+        y2 = min(orig_h, y2 + int(padding * height))
+        
+        cropped_face = image[y1:y2, x1:x2]
+        return cropped_face
+    
+    return None
+
 @app.post("/upload/")
 async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
     global model
@@ -42,13 +87,18 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Save the uploaded image temporarily to pass the file path to the model
+    # Detect and crop face
+    cropped_face = detect_and_crop_face(frame_bgr)
+    if cropped_face is None:
+        return {"error": "No face detected or alignment failed."}
+
+    # Save the cropped face temporarily
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-        cv2.imwrite(temp_file.name, frame_bgr)
+        cv2.imwrite(temp_file.name, cropped_face)
         temp_file_path = temp_file.name
 
     try:
-        # Process the uploaded image using the file path
+        # Process the cropped face 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.")