Backend Vision Update

backend/vision.py

@@ -0,0 +1,102 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO
+from scipy.spatial.distance import cdist
+from scipy.optimize import linear_sum_assignment
+
+# Load YOLOv8 nano model (downloads automatically if not found)
+# 'yolov8n.pt' is lightweight and fast for this purpose
+try:
+    model = YOLO("yolov8n.pt")
+except Exception as e:
+    print(f"Error loading YOLO model: {e}")
+    model = None
+
+# COCO Class IDs
+PERSON_CLASS_ID = 0
+CHAIR_CLASS_ID = 56
+
+def analyze_room_image(image_bytes: bytes):
+    if not model:
+        return {"error": "Vision model is not loaded"}
+
+    # Convert bytes to numpy array then to cv2 image
+    nparr = np.frombuffer(image_bytes, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    
+    if img is None:
+        return {"error": "Invalid image format"}
+
+    # Run inference
+    results = model(img)
+    
+    people = []
+    chairs = []
+
+    for result in results:
+        boxes = result.boxes
+        for box in boxes:
+            cls_id = int(box.cls[0])
+            conf = float(box.conf[0])
+            
+            # Extract center of bounding box
+            x1, y1, x2, y2 = box.xyxy[0]
+            cx = (x1 + x2) / 2.0
+            cy = (y1 + y2) / 2.0
+            centroid = [float(cx), float(cy)]
+
+            # Only consider detections with confidence > 0.3
+            if conf > 0.3:
+                if cls_id == PERSON_CLASS_ID:
+                    people.append({
+                        "centroid": centroid, 
+                        "box": [float(x1), float(y1), float(x2), float(y2)], 
+                        "conf": conf
+                    })
+                elif cls_id == CHAIR_CLASS_ID:
+                    chairs.append({
+                        "centroid": centroid, 
+                        "box": [float(x1), float(y1), float(x2), float(y2)], 
+                        "conf": conf
+                    })
+
+    num_people = len(people)
+    num_chairs = len(chairs)
+
+    pairs = []
+    
+    # Bipartite matching if both people and chairs exist
+    if num_people > 0 and num_chairs > 0:
+        people_coords = [p["centroid"] for p in people]
+        chairs_coords = [c["centroid"] for c in chairs]
+        
+        # Distance matrix (Euclidean distances)
+        dist_matrix = cdist(people_coords, chairs_coords, metric='euclidean')
+        
+        # Hungarian algorithm to minimize total distance for pairings
+        row_ind, col_ind = linear_sum_assignment(dist_matrix)
+        
+        for person_idx, chair_idx in zip(row_ind, col_ind):
+            distance = float(dist_matrix[person_idx, chair_idx])
+            pairs.append({
+                "person_index": int(person_idx),
+                "chair_index": int(chair_idx),
+                "distance": distance
+            })
+
+    # Basic logic: room is full if there are at least as many people as chairs. 
+    # Can be adjusted based on specific room definitions
+    is_full = num_people >= num_chairs if num_chairs > 0 else False
+
+    return {
+        "room_status": "full" if is_full else "available",
+        "counts": {
+            "people": num_people,
+            "chairs": num_chairs
+        },
+        "pairs": pairs,
+        "details": {
+            "people": people,
+            "chairs": chairs
+        }
+    }