Update app.py

app.py

@@ -1,22 +1,25 @@
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import Response
-from retinaface import RetinaFace
 from insightface.app import FaceAnalysis
 from rembg import remove
 from PIL import Image
 import numpy as np
 import io
 
+
 app = FastAPI()
 
+# Load face detection model (SCRFD)
 face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
-face_app.prepare(ctx_id=0)
+face_app.prepare(ctx_id=0, det_size=(640, 640))
+
 
 def resize_to_4x6(img):
     return img.resize((472, 709), Image.LANCZOS)
 
+
 @app.post("/process")
-async def process_image(file: UploadFile = File(...)):
+async def process_img(file: UploadFile = File(...)):
     img_bytes = await file.read()
     img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
     np_img = np.array(img)
@@ -24,18 +27,24 @@ async def process_image(file: UploadFile = File(...)):
     faces = face_app.get(np_img)
     if not faces:
         return {"error": "No face detected"}
+
+    # Extract aligned face from insightface
     face = faces[0]
     aligned = Image.fromarray(face.normed)
 
+    # Remove background
     no_bg = remove(aligned)
 
-    white_bg = Image.new("RGB", no_bg.size, (255, 255, 255))
-    white_bg.paste(no_bg, mask=no_bg.split()[3])
+    # Add white background
+    white = Image.new("RGB", no_bg.size, (255, 255, 255))
+    white.paste(no_bg, mask=no_bg.split()[3])
 
-    final_img = resize_to_4x6(white_bg)
+    # Resize to official 4x6 photo size
+    final_img = resize_to_4x6(white)
 
-    buffer = io.BytesIO()
-    final_img.save(buffer, format="JPEG")
-    buffer.seek(0)
+    # Return JPEG
+    buf = io.BytesIO()
+    final_img.save(buf, format="JPEG")
+    buf.seek(0)
 
-    return Response(buffer.getvalue(), media_type="image/jpeg")
+    return Response(buf.getvalue(), media_type="image/jpeg")