Update face_detection.py

face_detection.py

@@ -4,17 +4,29 @@ from PIL import Image, ImageOps
 import cv2
 
 # Load InsightFace models
-detector = insightface.model_zoo.get_model('retinaface_r50_v1')
-detector.prepare(ctx_id=-1)  # Use CPU, set ctx_id=0 for GPU
-
-landmark_model = insightface.model_zoo.get_model('2d106det')
-landmark_model.prepare(ctx_id=-1)  # Use CPU, set ctx_id=0 for GPU
+try:
+    detector = insightface.model_zoo.get_model('retinaface_r50_v1')
+    detector.prepare(ctx_id=-1)  # Use CPU, set ctx_id=0 for GPU
+except Exception as e:
+    print(f"Error loading face detection model: {e}")
+    detector = None
+
+try:
+    landmark_model = insightface.model_zoo.get_model('2d106det')
+    landmark_model.prepare(ctx_id=-1)  # Use CPU, set ctx_id=0 for GPU
+except Exception as e:
+    print(f"Error loading landmark model: {e}")
+    landmark_model = None
 
 def align(image_in, face_index=0, output_size=256):
+    if detector is None or landmark_model is None:
+        print("Models not loaded properly.")
+        return image_in, 0, None
+
     try:
         image_in = ImageOps.exif_transpose(image_in)
     except:
-        print("exif problem, not rotating")
+        print("Exif problem, not rotating")
 
     landmarks = list(get_landmarks(image_in))
     n_faces = len(landmarks)
@@ -29,6 +41,9 @@ def align(image_in, face_index=0, output_size=256):
 
 def get_landmarks(image):
     """Get landmarks from PIL image"""
+    if detector is None or landmark_model is None:
+        return []
+
     img = np.array(image)
     bboxes, _ = detector.detect(img, threshold=0.5, scale=1.0)
 
@@ -37,6 +52,8 @@ def get_landmarks(image):
         landmarks = landmark_model.get(img, bbox)
         yield landmarks
 
+# Remaining functions remain unchanged
+
 def image_align(src_img, face_landmarks, output_size=512, transform_size=2048, enable_padding=True, x_scale=1, y_scale=1, em_scale=0.1, alpha=False):
     # Align function modified from ffhq-dataset
     # See https://github.com/NVlabs/ffhq-dataset for license