Update app.py

app.py

@@ -26,9 +26,15 @@ from bson import ObjectId
 from bson.errors import InvalidId
 import httpx
 import uvicorn
-from PIL import Image
+from PIL import Image, ImageOps
 import io
 import requests
+try:
+    from pillow_heif import register_heif_opener
+except ImportError:
+    register_heif_opener = None
+else:
+    register_heif_opener()
 # DigitalOcean Spaces
 import boto3
 from botocore.client import Config
@@ -373,9 +379,7 @@ async def get_next_faceswap_sequence():
 
 
 def encode_image_as_jpeg(image_bytes: bytes) -> bytes:
-    img_bgr = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
-    if img_bgr is None:
-        raise HTTPException(400, "Invalid source image data")
+    img_bgr = decode_image_bytes(image_bytes, "source image")
 
     ok, jpeg_buffer = cv2.imencode(".jpg", img_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
     if not ok:
@@ -383,6 +387,20 @@ def encode_image_as_jpeg(image_bytes: bytes) -> bytes:
 
     return jpeg_buffer.tobytes()
 
+
+def decode_image_bytes(image_bytes: bytes, image_label: str = "image"):
+    img_bgr = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
+    if img_bgr is not None:
+        return img_bgr
+
+    try:
+        pil_img = Image.open(io.BytesIO(image_bytes))
+        pil_img = ImageOps.exif_transpose(pil_img).convert("RGB")
+    except Exception as exc:
+        raise HTTPException(400, f"Invalid {image_label} data") from exc
+
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
@@ -1006,8 +1024,8 @@ async def face_swap_api(
             response.raise_for_status()
             tgt_bytes = response.content
 
-        src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
-        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
+        src_bgr = decode_image_bytes(src_bytes, "source image")
+        tgt_bgr = decode_image_bytes(tgt_bytes, "target image")
 
         if src_bgr is None or tgt_bgr is None:
             raise HTTPException(400, "Invalid image data")
@@ -1020,7 +1038,7 @@ async def face_swap_api(
         # ------------------------------------------------------------------
         img2_rgb = None
         if img2_bytes:
-            img2_bgr = cv2.imdecode(np.frombuffer(img2_bytes, np.uint8), cv2.IMREAD_COLOR)
+            img2_bgr = decode_image_bytes(img2_bytes, "second source image")
             if img2_bgr is not None:
                 img2_rgb = cv2.cvtColor(img2_bgr, cv2.COLOR_BGR2RGB)
 
@@ -1201,8 +1219,8 @@ async def multi_face_swap_api(
         src_bytes = await source_image.read()
         tgt_bytes = await target_image.read()
 
-        src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
-        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
+        src_bgr = decode_image_bytes(src_bytes, "source image")
+        tgt_bgr = decode_image_bytes(tgt_bytes, "target image")
 
         if src_bgr is None or tgt_bgr is None:
             raise HTTPException(400, "Invalid image data")
@@ -1294,14 +1312,14 @@ async def face_swap_couple_api(
         # -----------------------------
         src_images = []
         img1_bytes = await image1.read()
-        src1 = cv2.imdecode(np.frombuffer(img1_bytes, np.uint8), cv2.IMREAD_COLOR)
+        src1 = decode_image_bytes(img1_bytes, "image1")
         if src1 is None:
             raise HTTPException(400, "Invalid image1 data")
         src_images.append(cv2.cvtColor(src1, cv2.COLOR_BGR2RGB))
 
         if image2:
             img2_bytes = await image2.read()
-            src2 = cv2.imdecode(np.frombuffer(img2_bytes, np.uint8), cv2.IMREAD_COLOR)
+            src2 = decode_image_bytes(img2_bytes, "image2")
             if src2 is not None:
                 src_images.append(cv2.cvtColor(src2, cv2.COLOR_BGR2RGB))
 
@@ -1507,7 +1525,7 @@ async def face_swap_couple_api(
             response.raise_for_status()
             tgt_bytes = response.content
 
-        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
+        tgt_bgr = decode_image_bytes(tgt_bytes, "target image")
         if tgt_bgr is None:
             raise HTTPException(400, "Invalid target image data")
 
@@ -3211,4 +3229,4 @@ if __name__ == "__main__":
 
 
 # if __name__ == "__main__":
-#     uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+#     uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)