Update app.py

app.py

@@ -5,6 +5,13 @@ from fastapi.responses import JSONResponse
 from ultralytics import YOLO
 import insightface
 
+from PIL import Image, ImageOps
+import pillow_heif
+import io
+
+# Enable HEIC/HEIF support
+pillow_heif.register_heif_opener()
+
 app = FastAPI()
 
 # ----------------------------
@@ -17,6 +24,10 @@ face_model = insightface.app.FaceAnalysis(name="buffalo_l")
 face_model.prepare(ctx_id=-1)
 
 
+# ----------------------------
+# Utility: Normalize embedding
+# ----------------------------
+
 def normalize(vec):
     vec = np.array(vec, dtype=np.float32)
     norm = np.linalg.norm(vec)
@@ -25,20 +36,52 @@ def normalize(vec):
     return (vec / norm).tolist()
 
 
+# ----------------------------
+# Decode Image (JPEG/PNG/WEBP/HEIC/HEIF)
+# ----------------------------
+
+def decode_image(body: bytes):
+    # 🔥 Fast path: OpenCV (JPEG/PNG/WebP)
+    np_arr = np.frombuffer(body, np.uint8)
+    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+    if image is not None:
+        return image
+
+    # 🔥 Fallback: Pillow (HEIC/HEIF support)
+    try:
+        image = Image.open(io.BytesIO(body))
+
+        # Auto-rotate based on EXIF
+        image = ImageOps.exif_transpose(image)
+
+        image = image.convert("RGB")
+        image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+        return image_np
+
+    except Exception:
+        return None
+
+
+# ----------------------------
+# Face Processing
+# ----------------------------
+
 def process_image_np(image_np):
     results = yolo(image_np)
     faces_output = []
 
     for r in results:
         boxes = r.boxes
+
         for box, cls, conf in zip(boxes.xyxy, boxes.cls, boxes.conf):
-            if int(cls) != 0:
+            if int(cls) != 0:  # YOLO class 0 = person
                 continue
+
             if float(conf) < 0.4:
                 continue
 
-            xmin, ymin, xmax, ymax = box.cpu().numpy()
-            xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])
+            xmin, ymin, xmax, ymax = map(int, box.cpu().numpy())
 
             h, w, _ = image_np.shape
             xmin = max(0, xmin)
@@ -54,6 +97,7 @@ def process_image_np(image_np):
 
             for face in detected_faces:
                 embedding = normalize(face.embedding)
+
                 fxmin, fymin, fxmax, fymax = face.bbox.astype(int)
 
                 faces_output.append({
@@ -72,15 +116,27 @@ def process_image_np(image_np):
     return faces_output
 
 
+# ----------------------------
+# API Endpoint
+# ----------------------------
+
 @app.post("/detect")
 async def detect(request: Request):
     body = await request.body()
 
-    np_arr = np.frombuffer(body, np.uint8)
-    image_np = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+    if not body:
+        return JSONResponse(
+            {"error": "Empty request body"},
+            status_code=400
+        )
+
+    image_np = decode_image(body)
 
     if image_np is None:
-        return JSONResponse({"error": "Invalid image"}, status_code=400)
+        return JSONResponse(
+            {"error": "Unsupported or invalid image format"},
+            status_code=400
+        )
 
     result = process_image_np(image_np)
     return result