Update main.py

main.py

@@ -36,19 +36,32 @@ def read_root():
 def health_check():
     return {"status": "ok"}
 
-
-
 @app.post("/extract", tags=["Extract Embeddings"])
 async def extract_embeddings(file: UploadFile = File(...)):
     # Load the image
     contents = await file.read()
     image = Image.open(io.BytesIO(contents)).convert('RGB')
 
-    # Preprocess the image
-    preprocessed_image = mtcnn(image)
-
+    # Detect faces
+    faces = mtcnn(image)
+    
+    # Check if any faces were detected
+    if faces is None:
+        return JSONResponse(content={"error": "No faces detected in the image"})
+    
+    # If faces is a list, take the first face. If it's a tensor, it's already the first (and only) face
+    if isinstance(faces, list):
+        face = faces[0]
+    else:
+        face = faces
+    
+    # Ensure the face tensor is 4D (batch_size, channels, height, width)
+    if face.dim() == 3:
+        face = face.unsqueeze(0)
+    
     # Extract the face embeddings
-    embeddings = resnet(preprocessed_image.unsqueeze(0)).detach().cpu().numpy().tolist()
+    with torch.no_grad():
+        embeddings = resnet(face).cpu().numpy().tolist()
 
     return JSONResponse(content={"embeddings": embeddings})