Spaces:
Running
Running
File size: 3,380 Bytes
9f43980 c64a277 9f43980 c64a277 45e0d59 c64a277 45e0d59 c64a277 9f43980 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import os
import cv2
import numpy as np
import logging
from ultralytics import YOLO
from pathlib import Path
logger = logging.getLogger(__name__)
class MultiTracker:
def __init__(self):
logger.info("Initializing Multi Tracker (Group Centroid)")
# Determine paths
# Get the directory containing the current script (multi_tracker.py)
base_dir = Path(__file__).parent.parent
# Model folder is now inside backend/
detector_model_path = base_dir / "Model" / "yolov8n-face.pt"
print(detector_model_path,"de")
try:
self.model = YOLO(detector_model_path)
logger.info(f"Loaded YOLO model from {detector_model_path}")
except Exception as e:
logger.error(f"Failed to load YOLO model: {e}")
self.model = None
def process_frame(self, frame):
"""
Process a single BGR image frame for group object tracking.
Returns a dictionary with tracking results (individual boxes + aggregate box).
"""
results_data = {
"individual_boxes": [],
"aggregate_box": None,
"centroid": None,
"error": None,
"frame_width": int(frame.shape[1]),
"frame_height": int(frame.shape[0])
}
if self.model is None:
results_data["error"] = "Model not initialized"
return results_data
try:
# RUN BYTETRACK (Detection + Tracking)
results = self.model.track(frame, persist=True, tracker="bytetrack.yaml", verbose=False)
if results and len(results) > 0 and results[0].boxes.id is not None:
boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
track_ids = results[0].boxes.id.cpu().numpy().astype(int)
all_x1, all_y1, all_x2, all_y2 = [], [], [], []
for box, track_id in zip(boxes, track_ids):
x1, y1, x2, y2 = box.tolist()
all_x1.append(x1)
all_y1.append(y1)
all_x2.append(x2)
all_y2.append(y2)
results_data["individual_boxes"].append({
"id": int(track_id),
"x1": int(x1), "y1": int(y1),
"x2": int(x2), "y2": int(y2)
})
# Calculate Aggregate Bounding Box if faces exist
if len(all_x1) > 0:
agg_x1 = int(min(all_x1))
agg_y1 = int(min(all_y1))
agg_x2 = int(max(all_x2))
agg_y2 = int(max(all_y2))
# Aggregate Centroid
agg_cx = (agg_x1 + agg_x2) // 2
agg_cy = (agg_y1 + agg_y2) // 2
results_data["aggregate_box"] = {
"x1": agg_x1, "y1": agg_y1,
"x2": agg_x2, "y2": agg_y2
}
results_data["centroid"] = {
"cx": agg_cx,
"cy": agg_cy
}
except Exception as e:
logger.error(f"Error during ByteTrack: {e}")
results_data["error"] = str(e)
return results_data
|