nedo-vision-worker-core 0.2.0__py3-none-any.whl

nedo_vision_worker_core/tracker/SFSORT.py

@@ -0,0 +1,325 @@
+
+import numpy as np
+use_lap=True
+try:
+    import lap
+except ImportError:
+    from scipy.optimize import linear_sum_assignment
+    use_lap=False
+ 
+# ******************************************************************** #
+# ***************************** Classes ****************************** #
+# ******************************************************************** # 
+class DotAccess(dict):
+    """Provides dot.notation access to dictionary attributes"""
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+    
+    
+class TrackState:
+    """Enumeration of possible states of a track"""    
+    Active = 0
+    Lost_Central = 1
+    Lost_Marginal = 2
+ 
+
+class Track:
+    """Handles basic track attributes and operations"""
+    
+    def __init__(self, bbox, frame_id, track_id):
+        """Track initialization"""                 
+        self.track_id = track_id
+        self.bbox = bbox 
+        self.state = TrackState.Active                      
+        self.last_frame = frame_id
+ 
+    def update(self, box, frame_id):
+        """Updates a matched track"""      
+        self.bbox = box
+        self.state = TrackState.Active
+        self.last_frame = frame_id
+
+class SFSORT:
+    """Multi-Object Tracking System"""
+    
+    def __init__(self, args):
+        """Initialize a tracker with given arguments""" 
+        args = DotAccess(args)     
+
+        # Register tracking arguments, setting default values if the argument is not provided      
+        if args.high_th is None:
+            self.high_th = 0.6
+        else:
+            self.high_th = self.clamp(args.high_th, 0, 1)
+            
+        if args.match_th_first is None:
+            self.match_th_first = 0.67
+        else:
+            self.match_th_first = self.clamp(args.match_th_first, 0, 0.67)  
+            
+        if args.new_track_th is None:
+            self.new_track_th = 0.7
+        else:
+            self.new_track_th = self.clamp(args.new_track_th, self.high_th, 1)          
+        
+        if args.low_th is None:
+            self.low_th = 0.1
+        else:
+            self.low_th = self.clamp(args.low_th, 0, self.high_th)       
+  
+        if args.match_th_second is None:
+            self.match_th_second = 0.3
+        else:
+            self.match_th_second = self.clamp(args.match_th_second, 0, 1)     
+  
+        self.dynamic_tuning = False
+        if args.dynamic_tuning is not None:
+            self.cth = 0.5
+            self.high_th_m = 0.0
+            self.new_track_th_m = 0.0
+            self.match_th_first_m = 0.0
+            if args.dynamic_tuning:
+                self.dynamic_tuning = True
+                if args.cth is not None:
+                    self.cth = self.clamp(args.cth, args.low_th, 1)      
+                if args.high_th_m is not None:
+                    self.high_th_m = self.clamp(args.high_th_m, 0.02, 0.1)   
+                if args.new_track_th_m is not None:
+                    self.new_track_th_m = self.clamp(args.new_track_th_m, 0.02, 0.08)                    
+                if args.match_th_first_m is not None:
+                    self.match_th_first_m = self.clamp(args.match_th_first_m, 0.02, 0.08)   
+                    
+        if args.marginal_timeout is None:
+            self.marginal_timeout = 0
+        else:
+            self.marginal_timeout = self.clamp(args.marginal_timeout, 0, 500) 
+            
+        if args.central_timeout is None:
+            self.central_timeout = 0
+        else:
+            self.central_timeout = self.clamp(args.central_timeout, 0, 1000)            
+        
+        self.l_margin = 0
+        self.r_margin = 0
+        if args.frame_width:
+            self.r_margin = args.frame_width
+            if args.horizontal_margin is not None:                             
+                self.l_margin = self.clamp(args.horizontal_margin, 0, args.frame_width) 
+                self.r_margin = self.clamp(args.frame_width - args.horizontal_margin, 0, args.frame_width) 
+        
+        self.t_margin = 0
+        self.b_margin = 0
+        if args.frame_height:
+            self.b_margin = args.frame_height
+            if args.vertical_margin is not None:                             
+                self.t_margin = self.clamp(args.vertical_margin, 0, args.frame_height) 
+                self.b_margin = self.clamp(args.frame_height - args.vertical_margin , 0, args.frame_height)   
+               
+        # Initialize the tracker
+        self.frame_no = 0      
+        self.id_counter = 0       
+        self.active_tracks = []         
+        self.lost_tracks = [] 
+        
+    def update(self, boxes, scores):
+        """Updates tracker with new detections"""
+        # Adjust dynamic arguments
+        hth = self.high_th 
+        nth = self.new_track_th 
+        mth = self.match_th_first     
+        if self.dynamic_tuning:
+            count = len(scores[scores>self.cth])     
+            if count < 1:
+              count = 1
+              
+            lnc = np.log10(count)           
+            hth = self.clamp(hth - (self.high_th_m * lnc), 0, 1)
+            nth = self.clamp(nth + (self.new_track_th_m * lnc), hth, 1)    
+            mth = self.clamp(mth - (self.match_th_first_m * lnc), 0, 0.67)  
+                  
+        # Increase frame number
+        self.frame_no += 1
+        
+        # Variable: Active tracks in the next frame
+        next_active_tracks = []
+        
+        # Remove long-time lost tracks      
+        all_lost_tracks = self.lost_tracks.copy()
+        for track in all_lost_tracks:
+            if track.state == TrackState.Lost_Central:
+                if self.frame_no - track.last_frame > self.central_timeout:
+                    self.lost_tracks.remove(track)                   
+            else:
+                if self.frame_no - track.last_frame > self.marginal_timeout:
+                    self.lost_tracks.remove(track)     
+                    
+        # Gather out all previous tracks
+        track_pool = self.active_tracks + self.lost_tracks  
+        
+        # Try to associate tracks with high score detections
+        unmatched_tracks = np.array([])
+        high_score = scores > hth 
+        if high_score.any():
+            definite_boxes = boxes[high_score]
+            definite_scores = scores[high_score]  
+            if track_pool:           
+                cost = self.calculate_cost(track_pool, definite_boxes) 
+                matches, unmatched_tracks, unmatched_detections = self.linear_assignment(cost, mth) 
+                # Update/Activate matched tracks
+                for track_idx, detection_idx in matches:
+                    box = definite_boxes[detection_idx]
+                    track = track_pool[track_idx]                   
+                    track.update(box, self.frame_no)
+                    next_active_tracks.append(track)
+                    # Remove re-identified tracks from lost list
+                    if track in self.lost_tracks:
+                        self.lost_tracks.remove(track)
+                # Identify eligible unmatched detections as new tracks
+                for detection_idx in unmatched_detections:                   
+                    if definite_scores[detection_idx] > nth:
+                        box = definite_boxes[detection_idx]
+                        track = Track(box, self.frame_no, self.id_counter)
+                        next_active_tracks.append(track)
+                        self.id_counter += 1                   
+            else:
+            	# Associate tracks of the first frame after object-free/null frames
+                for detection_idx, score in enumerate(definite_scores):                   
+                    if score > nth:
+                        box = definite_boxes[detection_idx]
+                        track = Track(box, self.frame_no, self.id_counter)
+                        next_active_tracks.append(track)
+                        self.id_counter += 1  
+        
+        # Add unmatched tracks to the lost list
+        unmatched_track_pool = []                                         
+        for track_address in unmatched_tracks:
+            unmatched_track_pool.append(track_pool[track_address])   
+        next_lost_tracks = unmatched_track_pool.copy()
+        
+        # Try to associate remained tracks with intermediate score detections 
+        intermediate_score = np.logical_and((self.low_th < scores), (scores < hth))      
+        if intermediate_score.any(): 
+            if len(unmatched_tracks):                 
+                possible_boxes = boxes[intermediate_score]
+                cost = self.calculate_cost(unmatched_track_pool, possible_boxes, iou_only=True)
+                matches, unmatched_tracks, unmatched_detections = self.linear_assignment(cost, self.match_th_second)
+                # Update/Activate matched tracks
+                for track_idx, detection_idx in matches:
+                    box = possible_boxes[detection_idx]
+                    track = unmatched_track_pool[track_idx]                  
+                    track.update(box, self.frame_no)
+                    next_active_tracks.append(track)
+                    # Remove re-identified tracks from lost list 
+                    if track in self.lost_tracks:
+                        self.lost_tracks.remove(track)
+                    next_lost_tracks.remove(track)                                  
+            
+        # All tracks are lost if there are no detections!
+        if not (high_score.any() or  intermediate_score.any()):
+            next_lost_tracks = track_pool.copy()
+            
+        # Update the list of lost tracks
+        for track in next_lost_tracks:
+            if track not in self.lost_tracks:
+                self.lost_tracks.append(track)
+                u = track.bbox[0] + (track.bbox[2] - track.bbox[0])/2             
+                v = track.bbox[1] + (track.bbox[3] - track.bbox[1])/2
+                if (self.l_margin < u < self.r_margin) and (self.t_margin < v < self.b_margin):                   
+                    track.state = TrackState.Lost_Central                        
+                else:
+                    track.state = TrackState.Lost_Marginal                      
+                           
+        # Update the list of active tracks
+        self.active_tracks = next_active_tracks.copy()
+
+        return np.asarray([[x.bbox, x.track_id] for x in next_active_tracks], dtype=object)
+
+    @staticmethod
+    def clamp(value, min_value, max_value):
+        """ Clamps a value within the specified minimum and maximum bounds."""
+        return max(min_value, min(value, max_value))
+
+    @staticmethod
+    def calculate_cost(tracks, boxes, iou_only=False):
+        """Calculates the association cost based on IoU and box similarity"""
+        eps = 1e-7
+        active_boxes = [track.bbox for track in tracks]   
+           
+        # Get the coordinates of bounding boxes
+        b1_x1, b1_y1, b1_x2, b1_y2 = np.array(active_boxes).T
+        b2_x1, b2_y1, b2_x2, b2_y2 = np.array(boxes).T      
+        
+        h_intersection = (np.minimum(b1_x2[:, None], b2_x2) - np.maximum(b1_x1[:, None], b2_x1)).clip(0)
+        w_intersection = (np.minimum(b1_y2[:, None], b2_y2) - np.maximum(b1_y1[:, None], b2_y1)).clip(0)
+        
+        # Calculate the intersection area
+        intersection =  h_intersection * w_intersection
+                     
+        # Calculate the union area
+        box1_height = b1_x2 - b1_x1
+        box2_height = b2_x2 - b2_x1
+        box1_width = b1_y2 - b1_y1   
+        box2_width = b2_y2 - b2_y1
+
+        box1_area = box1_height * box1_width
+        box2_area = box2_height * box2_width
+        
+        union = (box2_area + box1_area[:, None] - intersection + eps)
+        
+        # Calculate the IoU 
+        iou = intersection / union
+        
+        if iou_only:
+            return 1.0 - iou
+        
+        # Calculate the DIoU                    
+        centerx1 = (b1_x1 + b1_x2) / 2.0
+        centery1 = (b1_y1 + b1_y2) / 2.0
+        centerx2 = (b2_x1 + b2_x2) / 2.0
+        centery2 = (b2_y1 + b2_y2) / 2.0        
+        inner_diag = np.abs(centerx1[:, None] - centerx2) + np.abs(centery1[:, None] - centery2)
+           
+        xxc1 = np.minimum(b1_x1[:, None], b2_x1)
+        yyc1 = np.minimum(b1_y1[:, None], b2_y1)
+        xxc2 = np.maximum(b1_x2[:, None], b2_x2)
+        yyc2 = np.maximum(b1_y2[:, None], b2_y2)              
+        outer_diag = np.abs(xxc2 - xxc1) + np.abs(yyc2 - yyc1)
+       
+        diou = iou - (inner_diag / outer_diag)
+        
+        # Calculate the BBSI 
+        delta_w = np.abs(box2_width - box1_width[:, None])
+        sw = w_intersection / np.abs(w_intersection + delta_w + eps)
+        
+        delta_h = np.abs(box2_height - box1_height[:, None])     
+        sh = h_intersection / np.abs(h_intersection + delta_h + eps)
+               
+        bbsi = diou + sh + sw    
+        
+        # Normalize the BBSI
+        cost = (bbsi)/3.0
+
+        return 1.0 - cost 
+ 
+    @staticmethod
+    def linear_assignment(cost_matrix, thresh):
+        """Linear assignment"""
+        if cost_matrix.size == 0:
+            return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
+
+        if use_lap:
+            _, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
+            matches = [[ix, mx] for ix, mx in enumerate(x) if mx >= 0]
+            unmatched_a = np.where(x < 0)[0]
+            unmatched_b = np.where(y < 0)[0]
+        else:
+            row_ind, col_ind = linear_sum_assignment(cost_matrix)  
+            matches = np.array([[row, col] for row, col in zip(row_ind, col_ind) if cost_matrix[row, col] <= thresh])  
+            matched_rows = set(row_ind)  
+            matched_cols = set(col_ind)  
+            unmatched_a = np.array([i for i in range(cost_matrix.shape[0]) if i not in matched_rows])  
+            unmatched_b = np.array([j for j in range(cost_matrix.shape[1]) if j not in matched_cols])  
+            
+        return matches, unmatched_a, unmatched_b
+    

nedo_vision_worker_core/tracker/TrackerManager.py

@@ -0,0 +1,163 @@
+import uuid
+import time
+import numpy as np
+from .SFSORT import SFSORT
+
+class TrackerManager:
+    def __init__(self, attribute_labels=None, exclusive_attribute_groups=None):
+        self.tracker = SFSORT({
+            "dynamic_tuning": True,
+            "cth": 0.5,
+            "high_th": 0.6,
+            "match_th_first": 0.67,
+            "match_th_second": 0.2,
+            "low_th": 0.1,
+            "new_track_th": 0.7,
+            "marginal_timeout": 7,
+            "central_timeout": 30
+        })
+        self.track_uuid_map = {}
+        self.track_count_map = {}
+        self.track_attributes_presence = {}
+        self.track_last_seen = {}
+        self.track_timeout_seconds = 5
+        self.attribute_labels = attribute_labels or []
+        self.exclusive_attribute_groups = exclusive_attribute_groups or []
+
+    def track_objects(self, detections):
+        if not detections:
+            self._cleanup_stale_tracks()
+            return []
+
+        bboxes = np.array([d["bbox"] for d in detections], dtype=np.float32)
+        confidences = np.array([d["confidence"] for d in detections], dtype=np.float32)
+        tracks = self.tracker.update(bboxes, confidences)
+
+        results = self._generate_tracking_results(detections, tracks)
+        self._cleanup_stale_tracks()
+        return results
+
+    def _generate_tracking_results(self, detections, tracks):
+        tracked_results = []
+        detection_map = {tuple(d["bbox"]): d for d in detections}
+
+        for track in tracks:
+            track_id = int(track[1])
+            bbox = track[0].tolist()
+            data = detection_map.get(tuple(bbox))
+
+            if not data:
+                continue
+
+            obj_uuid = self._assign_uuid(track_id)
+
+            self.track_count_map[obj_uuid] += 1
+            self.track_last_seen[obj_uuid] = time.time()  # Time-based last seen
+
+            attributes = data.get("attributes", [])
+            filtered_attributes = self._filter_exclusive_attributes(attributes)
+            self._update_attribute_presence(obj_uuid, filtered_attributes)
+
+            for attr in filtered_attributes:
+                label = attr["label"]
+                if label in self.track_attributes_presence[obj_uuid]:
+                    attr["count"] = self.track_attributes_presence[obj_uuid][label]
+
+            tracked_results.append({
+                "uuid": obj_uuid,
+                "track_id": track_id,
+                "detections": self.track_count_map[obj_uuid],
+                "label": data["label"],
+                "confidence": data["confidence"],
+                "bbox": bbox,
+                "attributes": filtered_attributes
+            })
+
+        return tracked_results
+
+    def _assign_uuid(self, track_id):
+        if track_id not in self.track_uuid_map:
+            new_uuid = str(uuid.uuid4())
+            self.track_uuid_map[track_id] = new_uuid
+            self.track_count_map[new_uuid] = 0
+            self.track_attributes_presence[new_uuid] = {attr: 0 for attr in self.attribute_labels}
+        return self.track_uuid_map[track_id]
+
+    def _filter_exclusive_attributes(self, attributes):
+        if not attributes:
+            return []
+
+        # Group attributes by label to handle multiple instances
+        attrs_by_label = {}
+        for attr in attributes:
+            label = attr["label"]
+            if label not in attrs_by_label:
+                attrs_by_label[label] = []
+            attrs_by_label[label].append(attr)
+
+        # Multi-instance classes that can have multiple instances per person
+        MULTI_INSTANCE_CLASSES = ["boots", "gloves", "goggles", "no_gloves"]
+        # Negative classes that are exclusive with their positive counterparts
+        NEGATIVE_CLASSES = ["no_helmet", "no_vest", "no_goggles", "no_boots"]
+        
+        # For exclusive groups, keep only the highest confidence per group
+        filtered_attrs = []
+        for group in self.exclusive_attribute_groups:
+            group_attrs = []
+            for label in group:
+                if label in attrs_by_label:
+                    group_attrs.extend(attrs_by_label[label])
+            if group_attrs:
+                # Keep only the highest confidence in this exclusive group
+                best = max(group_attrs, key=lambda a: a["confidence"])
+                filtered_attrs.append(best)
+
+        # For multi-instance classes, add all instances (but respect exclusive logic above)
+        exclusive_labels = set(l for group in self.exclusive_attribute_groups for l in group)
+        for label, attrs in attrs_by_label.items():
+            # Skip if already handled by exclusive logic
+            if label in exclusive_labels:
+                continue
+            # Add all instances of multi-instance classes
+            if label in MULTI_INSTANCE_CLASSES:
+                filtered_attrs.extend(attrs)
+        
+        # Special case: for multi-instance classes, we want to allow multiple instances
+        # even if they were in exclusive groups, but we need to handle the negative classes properly
+        for label, attrs in attrs_by_label.items():
+            # If this is a multi-instance class and not a negative class, add all instances
+            if label in MULTI_INSTANCE_CLASSES and label not in NEGATIVE_CLASSES:
+                # Check which attributes were already added
+                for attr in attrs:
+                    already_added = any(
+                        a["label"] == label and list(a["bbox"]) == list(attr["bbox"]) 
+                        for a in filtered_attrs
+                    )
+                    if not already_added:
+                        filtered_attrs.append(attr)
+
+        return filtered_attrs
+
+    def _update_attribute_presence(self, uuid, attributes):
+        current_frame_attrs = set(attr["label"] for attr in attributes)
+        for label in self.attribute_labels:
+            if label in current_frame_attrs:
+                self.track_attributes_presence[uuid][label] += 1
+            else:
+                self.track_attributes_presence[uuid][label] = 0
+
+    def _cleanup_stale_tracks(self):
+        now = time.time()
+        expired = [
+            uuid for uuid, last_seen in self.track_last_seen.items()
+            if now - last_seen > self.track_timeout_seconds
+        ]
+
+        for obj_uuid in expired:
+            track_ids_to_remove = [tid for tid, uid in self.track_uuid_map.items() if uid == obj_uuid]
+            for tid in track_ids_to_remove:
+                self.track_uuid_map.pop(tid, None)
+
+            self.track_count_map.pop(obj_uuid, None)
+            self.track_attributes_presence.pop(obj_uuid, None)
+            self.track_last_seen.pop(obj_uuid, None)

nedo_vision_worker_core/tracker/__init__.py

@@ -0,0 +1 @@
+ 

nedo_vision_worker_core/util/BoundingBoxMetrics.py

@@ -0,0 +1,53 @@
+class BoundingBoxMetrics:
+    """Computes IoU and coverage for bounding boxes."""
+
+    @staticmethod
+    def compute_iou(box1, box2):
+        """
+        Computes Intersection over Union (IoU) between two bounding boxes.
+        Args:
+            box1, box2: [x1, y1, x2, y2] (coordinates of two bounding boxes)
+        Returns:
+            IoU score (float between 0 and 1)
+        """
+        x1, y1, x2, y2 = box1
+        x1_p, y1_p, x2_p, y2_p = box2
+
+        # Compute intersection
+        inter_x1 = max(x1, x1_p)
+        inter_y1 = max(y1, y1_p)
+        inter_x2 = min(x2, x2_p)
+        inter_y2 = min(y2, y2_p)
+
+        inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
+
+        # Compute union
+        box1_area = (x2 - x1) * (y2 - y1)
+        box2_area = (x2_p - x1_p) * (y2_p - y1_p)
+        union_area = box1_area + box2_area - inter_area
+
+        return inter_area / union_area if union_area > 0 else 0
+
+    @staticmethod
+    def compute_coverage(box1, box2):
+        """
+        Computes the coverage percentage of box2 inside box1.
+        Coverage is defined as the intersection area over box2's area.
+        Args:
+            box1, box2: [x1, y1, x2, y2] (coordinates of two bounding boxes)
+        Returns:
+            Coverage ratio (float between 0 and 1)
+        """
+        x1, y1, x2, y2 = box1
+        x1_o, y1_o, x2_o, y2_o = box2
+
+        # Compute intersection
+        inter_x1 = max(x1, x1_o)
+        inter_y1 = max(y1, y1_o)
+        inter_x2 = min(x2, x2_o)
+        inter_y2 = min(y2, y2_o)
+
+        inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
+        box2_area = (x2_o - x1_o) * (y2_o - y1_o)
+
+        return inter_area / box2_area if box2_area > 0 else 0