Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile, Form
+from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import StreamingResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
 import torch
@@ -13,6 +13,7 @@ app = FastAPI()
 
 # Load model and necessary components
 model = None
+face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
 
 def load_model():
     global model
@@ -24,7 +25,7 @@ def load_model():
 logging.basicConfig(level=logging.INFO)
 
 @app.post("/upload/")
-async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
+async def process_image(file: UploadFile = File(...)):
     global model
     if model is None:
         load_model()
@@ -42,6 +43,18 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
+    # Detect faces in the image
+    gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+
+    if len(faces) == 0:
+        logging.error("No faces detected in the image.")
+        return {"error": "No faces detected in the image."}
+
+    # Use the first detected face
+    (x, y, w, h) = faces[0]
+    top, bottom, left, right = y, y + h, x, x + w
+
     # Save the uploaded image temporarily to pass the file path to the model
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
         cv2.imwrite(temp_file.name, frame_bgr)