Update app.py

app.py

@@ -6,27 +6,27 @@ from ultralytics import YOLO
 from fastapi import FastAPI, File, UploadFile
 import uvicorn
 
-# Load YOLO model with optimized settings
+# Load YOLO model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = YOLO("yolov11s-face.pt").to(device)
 
 # Constants for distance estimation
 KNOWN_DISTANCE = 50  # cm
 KNOWN_FACE_WIDTH = 14  # cm
-REF_IMAGE_FACE_WIDTH = 120  # Reference face width in pixels at the known distance
+REF_IMAGE_FACE_WIDTH = 120  # Reference face width in pixels
 FOCAL_LENGTH = (REF_IMAGE_FACE_WIDTH * KNOWN_DISTANCE) / KNOWN_FACE_WIDTH
 SCALING_FACTOR = 2.0  # Adjust based on testing
 
 # FastAPI initialization
 app = FastAPI()
 
-# Optimized tracking of previous detections using a deque
+# Tracking history of detections
 MAX_HISTORY = 10
 detected_people_history = deque(maxlen=MAX_HISTORY)
 
 @app.post("/detect")
 async def detect_faces(file: UploadFile = File(...)):
-    # Read the image and decode it efficiently
+    # Read the image and decode it
     contents = await file.read()
     image_np = np.frombuffer(contents, np.uint8)
     frame = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
@@ -41,26 +41,29 @@ async def detect_faces(file: UploadFile = File(...)):
     new_people_data = {}
     person_id = 1
     frame_width = resized_frame.shape[1]
-    frame_center = frame_width // 2  # Calculate the frame's center
 
-    center_threshold = frame_width * 0.15  # 15% margin for center detection
+    # Define Left, Center, Right based on frame width
+    left_count, center_count, right_count = 0, 0, 0
 
     for result in results:
-        for box in result.boxes.data.tolist():  # Convert tensor to list
-            x1, y1, x2, y2, conf, _ = box[:6]  # Extract values
-            x1, y1, x2, y2 = map(int, [x1 * 2, y1 * 2, x2 * 2, y2 * 2])  # Scale back
+        for box in result.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            conf = box.conf[0].item()
 
-            if conf > 0.5:
+            if conf > 0.5:  # Confidence threshold
                 center_x = (x1 + x2) // 2
-                face_width_pixels = x2 - x1
+                face_width_pixels = x2 - x1  # Detected face width in pixels
 
-                # **Improved Position Logic**
-                if center_x < frame_center - center_threshold:
+                # **Updated Position Logic**
+                if center_x < frame_width // 3:
                     position = "Left"
-                elif center_x > frame_center + center_threshold:
+                    left_count += 1
+                elif center_x > 2 * frame_width // 3:
                     position = "Right"
+                    right_count += 1
                 else:
-                    position = "Center"  # Corrected condition for center
+                    position = "Center"
+                    center_count += 1
 
                 # Calculate distance
                 estimated_distance = (
@@ -77,6 +80,6 @@ async def detect_faces(file: UploadFile = File(...)):
     # Return only if there's a change
     if not detected_people_history or new_people_data != detected_people_history[-1]:
         detected_people_history.append(new_people_data)
-        return {"people": new_people_data}
+        return {"people": new_people_data, "counts": {"Left": left_count, "Center": center_count, "Right": right_count}}
 
     return {"people": []}  # No change, return empty response