Update app.py

app.py

@@ -4,53 +4,64 @@ import numpy as np
 from fastapi import FastAPI, File, UploadFile
 from ultralytics import YOLO
 from typing import List, Dict
-import uvicorn
+from io import BytesIO
 
 app = FastAPI()
 
-# Load YOLO model
-model = YOLO("yolov11s-face.pt")  # Change model path as needed
+# Load YOLO Model
+model = YOLO("yolov11s-face.pt")
 
-# Constants for Distance Estimation
+# Constants for Distance Calculation
 KNOWN_DISTANCE = 50  # cm
 KNOWN_FACE_WIDTH = 14  # cm
-REF_IMAGE_FACE_WIDTH = 120  # px (Change this based on calibration)
+REF_IMAGE_FACE_WIDTH = 120  # pixels
 FOCAL_LENGTH = (REF_IMAGE_FACE_WIDTH * KNOWN_DISTANCE) / KNOWN_FACE_WIDTH
-SCALING_FACTOR = 2.0  # Adjust based on testing
+SCALING_FACTOR = 2.0  # Adjust based on real-world testing
 
-@app.post("/detect")
-async def detect_faces(file: UploadFile = File(...)) -> Dict[str, List[Dict[str, str]]]:
-    # Read image from file
-    contents = await file.read()
-    np_arr = np.frombuffer(contents, np.uint8)
-    frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
-
-    frame_width = frame.shape[1]
-
-    # Run YOLO detection
-    results = model(frame)
-
-    faces_detected = []
-
-    for i, result in enumerate(results):
-        for j, box in enumerate(result.boxes):
+# Function to Process Frame & Detect Faces
+def process_frame(image: np.ndarray) -> List[Dict]:
+    results = model(image)
+    frame_width = image.shape[1]
+    people = []
+    
+    for idx, result in enumerate(results):
+        for i, box in enumerate(result.boxes):
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             conf = box.conf[0].item()
-
+            
             if conf > 0.5:  # Confidence threshold
                 center_x = (x1 + x2) // 2
-                face_width_pixels = x2 - x1  # Detected face width in pixels
-
-                # Calculate distance
+                face_width_pixels = x2 - x1
+                
+                # Determine Position
+                if center_x < frame_width // 3:
+                    position = "Left"
+                elif center_x > 2 * frame_width // 3:
+                    position = "Right"
+                else:
+                    position = "Center"
+                
+                # Estimate Distance
                 if face_width_pixels > 0:
                     estimated_distance = (FOCAL_LENGTH * KNOWN_FACE_WIDTH) / face_width_pixels
-                    estimated_distance *= SCALING_FACTOR  # Apply scaling fix
+                    estimated_distance *= SCALING_FACTOR
                 else:
                     estimated_distance = -1  # Error case
+                
+                people.append({
+                    f"person{i+1}": {
+                        "distance": round(estimated_distance, 1),
+                        "position": position
+                    }
+                })
+    
+    return people
 
-                person_data = {
-                    f"person{j+1}": f"{estimated_distance:.1f} cm"
-                }
-                faces_detected.append(person_data)
-
-    return {"faces": faces_detected}
+@app.post("/detect")
+async def detect_faces(file: UploadFile = File(...)):
+    image_data = await file.read()
+    nparr = np.frombuffer(image_data, np.uint8)
+    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    
+    faces = process_frame(image)
+    return {"people": faces}