Update main.py

main.py

@@ -1,74 +1,67 @@
-from fastapi import FastAPI, HTTPException, UploadFile, File
-from fastapi.responses import FileResponse
-from insightface.app import FaceAnalysis
+
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
 import insightface
+from insightface.app import FaceAnalysis
+from insightface.model_zoo import get_model
+import numpy as np
 import cv2
-import shutil
-from tempfile import NamedTemporaryFile
-
-assert insightface.__version__ >= '0.7'
+import io
 
 app = FastAPI()
 
-def prepare_app():
-    app = FaceAnalysis(name='buffalo_l')
-    app.prepare(ctx_id=0, det_size=(640, 640))
-    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)
-    return app, swapper
+# Initialize FaceAnalysis app and swapper
+face_analysis = FaceAnalysis(name='buffalo_l')
+face_analysis.prepare(ctx_id=0, det_size=(640, 640))
+swapper = get_model('inswapper_128.onnx', download=True, download_zip=True)
+
+@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."""
+    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))
+    if not faces_source:
+        raise HTTPException(status_code=400, detail="No faces detected in the source image.")
+    source_face = get_face(faces_source, source_face_index)
+
+    faces_destination = sort_faces(face_analysis.get(destination_image))
+    if not faces_destination:
+        raise HTTPException(status_code=400, detail="No faces detected in the destination image.")
+    destination_face = get_face(faces_destination, destination_face_index)
 
-face_app, swapper = prepare_app()
+    # 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")
 
 def sort_faces(faces):
     return sorted(faces, key=lambda x: x.bbox[0])
 
 def get_face(faces, face_id):
     if len(faces) < face_id or face_id < 1:
-        raise HTTPException(status_code=400, detail=f"The image includes only {len(faces)} faces, however, you asked for face {face_id}")
-    return faces[face_id-1]
-
-def swap_faces(source_image_path, source_face_index, destination_image_path, destination_face_index):
-    source_image = cv2.imread(source_image_path)
-    destination_image = cv2.imread(destination_image_path)
-
-    if source_image is None:
-        raise HTTPException(status_code=400, detail="Source image could not be read.")
-    if destination_image is None:
-        raise HTTPException(status_code=400, detail="Destination image could not be read.")
+        raise ValueError(f"The image includes only {len(faces)} faces, however, you asked for face {face_id}")
+    return faces[face_id - 1]
 
-    faces = sort_faces(face_app.get(source_image))
-    if not faces:
-        raise HTTPException(status_code=400, detail="No faces detected in the source image.")
-    source_face = get_face(faces, source_face_index)
-    
-    res_faces = sort_faces(face_app.get(destination_image))
-    if not res_faces:
-        raise HTTPException(status_code=400, detail="No faces detected in the destination image.")
-    res_face = get_face(res_faces, destination_face_index)
+@app.exception_handler(ValueError)
+async def value_error_handler(request, exc):
+    """Custom exception handler to return JSON error responses for ValueError."""
+    return JSONResponse(status_code=400, content={"detail": str(exc)})
 
-    result = swapper.get(destination_image, res_face, source_face, paste_back=True)
-    return result
-
-@app.post("/swap_faces/")
-async def swap_faces_endpoint(
-    source_file: UploadFile = File(...),
-    source_face_index: int = 1,
-    destination_file: UploadFile = File(...),
-    destination_face_index: int = 1
-):
-    with NamedTemporaryFile(delete=False, suffix=".jpg") as source_temp_file:
-        shutil.copyfileobj(source_file.file, source_temp_file)
-        source_image_path = source_temp_file.name
 
-    with NamedTemporaryFile(delete=False, suffix=".jpg") as destination_temp_file:
-        shutil.copyfileobj(destination_file.file, destination_temp_file)
-        destination_image_path = destination_temp_file.name
 
-    try:
-        result = swap_faces(source_image_path, source_face_index, destination_image_path, destination_face_index)
-        result_path = "result.jpg"
-        cv2.imwrite(result_path, result)
-        return FileResponse(result_path, media_type="image/jpeg")
-    except HTTPException as e:
-        raise e
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))