Update main.py

main.py

@@ -26,37 +26,54 @@ from fastapi.responses import StreamingResponse
 
 
 from fastapi import Form
+from fastapi import HTTPException
 
 @app.post("/swap_faces/")
 async def swap_faces(source_file: UploadFile = File(...),
                      source_face_index: int = Form(...),
                      destination_file: UploadFile = File(...),
                      destination_face_index: int = Form(...)):
-    """Swaps faces between the source and destination images based on the specified face indices."""
-    # Your existing code for swapping faces...
-
-    source_bytes = await source_file.read()
-    destination_bytes = await destination_file.read()
-
-    # Decode images
-    source_image = cv2.imdecode(np.frombuffer(source_bytes, np.uint8), cv2.IMREAD_COLOR)
-    destination_image = cv2.imdecode(np.frombuffer(destination_bytes, np.uint8), cv2.IMREAD_COLOR)
-
-    # Face detection and sorting
-    faces_source = sort_faces(face_analysis.get(source_image))
-    source_face = get_face(faces_source, source_face_index)
-
-    faces_destination = sort_faces(face_analysis.get(destination_image))
-    destination_face = get_face(faces_destination, destination_face_index)
-
-    # Swap faces
-    result_image = swapper.get(destination_image, destination_face, source_face, paste_back=True)
-    
-    # Convert result_image back to bytes
-    _, result_bytes = cv2.imencode('.jpg', result_image)
-    
-    # Return the image bytes as a streaming response
-    return StreamingResponse(io.BytesIO(result_bytes), media_type="image/jpeg")
+    try:
+        # Your existing code for swapping faces...
+
+        source_bytes = await source_file.read()
+        destination_bytes = await destination_file.read()
+
+        # Decode images
+        source_image = cv2.imdecode(np.frombuffer(source_bytes, np.uint8), cv2.IMREAD_COLOR)
+        destination_image = cv2.imdecode(np.frombuffer(destination_bytes, np.uint8), cv2.IMREAD_COLOR)
+
+        # Face detection and sorting
+        faces_source = sort_faces(face_analysis.get(source_image))
+        source_face = get_face(faces_source, source_face_index)
+
+        faces_destination = sort_faces(face_analysis.get(destination_image))
+        destination_face = get_face(faces_destination, destination_face_index)
+
+        # Check if there are faces in both images
+        if not faces_source:
+            raise ValueError("Source image has no faces")
+        if not faces_destination:
+            raise ValueError("Destination image has no faces")
+
+        # Swap faces
+        result_image = swapper.get(destination_image, destination_face, source_face, paste_back=True)
+        
+        # Convert result_image back to bytes
+        _, result_bytes = cv2.imencode('.jpg', result_image)
+        
+        # Return the image bytes as a streaming response
+        return StreamingResponse(io.BytesIO(result_bytes), media_type="image/jpeg")
+    except ValueError as ve:
+        # Log the error message to console
+        print(f"An error occurred: {ve}")
+        # Return the error message as part of the response
+        return {"error": str(ve)}
+    except Exception as e:
+        # Log other types of errors to console
+        print(f"An error occurred: {e}")
+        # Raise HTTPException with 500 Internal Server Error status code and a generic error message
+        raise HTTPException(status_code=500, detail="An unexpected error occurred")