bplusplus 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

bplusplus/{hierarchical/test.py → test.py}

@@ -23,7 +23,7 @@ import sys
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-def test_multitask(species_list, test_set, yolo_weights, hierarchical_weights, output_dir="."):
+def test(species_list, test_set, yolo_weights, hierarchical_weights, output_dir="."):
     """
     Run the two-stage classifier on a test set.
     
@@ -115,17 +115,17 @@ class HierarchicalInsectClassifier(nn.Module):
 def get_taxonomy(species_list):
     """
     Retrieves taxonomic information for a list of species from GBIF API.
-    Creates a hierarchical taxonomy dictionary with order, family, and species relationships.
+    Creates a hierarchical taxonomy dictionary with family, genus, and species relationships.
     """
     taxonomy = {1: [], 2: {}, 3: {}}
-    species_to_family = {}
-    family_to_order = {}
+    species_to_genus = {}
+    genus_to_family = {}
     
     logger.info(f"Building taxonomy from GBIF for {len(species_list)} species")
     
     print("\nTaxonomy Results:")
     print("-" * 80)
-    print(f"{'Species':<30} {'Order':<20} {'Family':<20} {'Status'}")
+    print(f"{'Species':<30} {'Family':<20} {'Genus':<20} {'Status'}")
     print("-" * 80)
     
     for species_name in species_list:
@@ -136,23 +136,23 @@ def get_taxonomy(species_list):
             
             if data.get('status') == 'ACCEPTED' or data.get('status') == 'SYNONYM':
                 family = data.get('family')
-                order = data.get('order')
+                genus = data.get('genus')
                 
-                if family and order:
+                if family and genus:
                     status = "OK"
                     
-                    print(f"{species_name:<30} {order:<20} {family:<20} {status}")
+                    print(f"{species_name:<30} {family:<20} {genus:<20} {status}")
                     
-                    species_to_family[species_name] = family
-                    family_to_order[family] = order
+                    species_to_genus[species_name] = genus
+                    genus_to_family[genus] = family
                     
-                    if order not in taxonomy[1]:
-                        taxonomy[1].append(order)
+                    if family not in taxonomy[1]:
+                        taxonomy[1].append(family)
                     
-                    taxonomy[2][family] = order
-                    taxonomy[3][species_name] = family
+                    taxonomy[2][genus] = family
+                    taxonomy[3][species_name] = genus
                 else:
-                    error_msg = f"Species '{species_name}' found in GBIF but family and order not found, could be spelling error in species, check GBIF"
+                    error_msg = f"Species '{species_name}' found in GBIF but family and genus not found, could be spelling error in species, check GBIF"
                     logger.error(error_msg)
                     print(f"{species_name:<30} {'Not found':<20} {'Not found':<20} ERROR")
                     print(f"Error: {error_msg}")
@@ -174,24 +174,24 @@ def get_taxonomy(species_list):
     taxonomy[1] = sorted(list(set(taxonomy[1])))
     print("-" * 80)
     
-    num_orders = len(taxonomy[1])
-    num_families = len(taxonomy[2])
+    num_families = len(taxonomy[1])
+    num_genera = len(taxonomy[2])
     num_species = len(taxonomy[3])
     
-    print("\nOrder indices:")
-    for i, order in enumerate(taxonomy[1]):
-        print(f"  {i}: {order}")
-    
     print("\nFamily indices:")
-    for i, family in enumerate(taxonomy[2].keys()):
+    for i, family in enumerate(taxonomy[1]):
         print(f"  {i}: {family}")
     
+    print("\nGenus indices:")
+    for i, genus in enumerate(taxonomy[2].keys()):
+        print(f"  {i}: {genus}")
+    
     print("\nSpecies indices:")
     for i, species in enumerate(species_list):
         print(f"  {i}: {species}")
     
-    logger.info(f"Taxonomy built: {num_orders} orders, {num_families} families, {num_species} species")
-    return taxonomy, species_to_family, family_to_order
+    logger.info(f"Taxonomy built: {num_families} families, {num_genera} genera, {num_species} species")
+    return taxonomy, species_to_genus, genus_to_family
 
 def create_mappings(taxonomy):
     """Create index mappings from taxonomy"""
@@ -243,13 +243,20 @@ class TestTwoStage:
                 if "species_list" in checkpoint:
                     saved_species = checkpoint["species_list"]
                     print(f"Saved model was trained on: {', '.join(saved_species)}")
-    
-                taxonomy, species_to_family, family_to_order = get_taxonomy(species_names)
+                
+                # Use saved taxonomy mappings if available
+                if "species_to_genus" in checkpoint and "genus_to_family" in checkpoint:
+                    species_to_genus = checkpoint["species_to_genus"]
+                    genus_to_family = checkpoint["genus_to_family"]
+                else:
+                    # Fallback: fetch from GBIF but this may cause index mismatches
+                    print("Warning: No taxonomy mappings in checkpoint, fetching from GBIF")
+                    _, species_to_genus, genus_to_family = get_taxonomy(species_names)
             else:
-                taxonomy, species_to_family, family_to_order = get_taxonomy(species_names)
+                taxonomy, species_to_genus, genus_to_family = get_taxonomy(species_names)
         else:
             state_dict = checkpoint
-            taxonomy, species_to_family, family_to_order = get_taxonomy(species_names)
+            taxonomy, species_to_genus, genus_to_family = get_taxonomy(species_names)
         
         level_to_idx, idx_to_level = create_mappings(taxonomy)
         
@@ -259,8 +266,6 @@ class TestTwoStage:
         if hasattr(taxonomy, "items"):
             num_classes_per_level = [len(classes) if isinstance(classes, list) else len(classes.keys()) 
                                     for level, classes in taxonomy.items()]
-        else:
-            num_classes_per_level = [4, 5, 9]  # Example values, adjust as needed
         
         print(f"Using model with class counts: {num_classes_per_level}")
         
@@ -287,8 +292,6 @@ class TestTwoStage:
         self.classification_model.eval()
 
         self.classification_transform = transforms.Compose([
-            transforms.Resize((768, 768)),  # Fixed size for all validation images
-            transforms.CenterCrop(640),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
@@ -296,8 +299,8 @@ class TestTwoStage:
         print("Model successfully loaded")
         print(f"Using species: {', '.join(species_names)}")
         
-        self.species_to_family = species_to_family
-        self.family_to_order = family_to_order
+        self.species_to_genus = species_to_genus
+        self.genus_to_family = genus_to_family
 
     def get_frames(self, test_dir):
         image_dir = os.path.join(test_dir, "images")
@@ -305,10 +308,10 @@ class TestTwoStage:
         
         predicted_frames = []
         predicted_family_frames = []
-        predicted_order_frames = []
+        predicted_genus_frames = []
         true_species_frames = []
         true_family_frames = []
-        true_order_frames = []
+        true_genus_frames = []
         image_names = []
 
         start_time = time.time()  # Start timing
@@ -326,7 +329,7 @@ class TestTwoStage:
             detections = results[0].boxes
             predicted_frame = []
             predicted_family_frame = []
-            predicted_order_frame = []
+            predicted_genus_frame = []
 
             if detections:
                 for box in detections:
@@ -346,13 +349,13 @@ class TestTwoStage:
                         outputs = self.classification_model(input_tensor)
                     
                     # Get all taxonomic level predictions
-                    order_output = outputs[0]    # First output is order (level 1)
-                    family_output = outputs[1]   # Second output is family (level 2)
+                    family_output = outputs[0]   # First output is family (level 1)
+                    genus_output = outputs[1]    # Second output is genus (level 2)
                     species_output = outputs[2]  # Third output is species (level 3)
                     
                     # Get prediction indices
-                    order_idx = order_output.argmax(dim=1).item()
                     family_idx = family_output.argmax(dim=1).item()
+                    genus_idx = genus_output.argmax(dim=1).item()
                     species_idx = species_output.argmax(dim=1).item()
                     
                     img_height, img_width, _ = frame.shape
@@ -367,15 +370,15 @@ class TestTwoStage:
                     # Add predictions for each taxonomic level
                     predicted_frame.append([species_idx] + box_coords)
                     predicted_family_frame.append([family_idx] + box_coords)
-                    predicted_order_frame.append([order_idx] + box_coords)
+                    predicted_genus_frame.append([genus_idx] + box_coords)
 
             predicted_frames.append(predicted_frame if predicted_frame else [])
             predicted_family_frames.append(predicted_family_frame if predicted_family_frame else [])
-            predicted_order_frames.append(predicted_order_frame if predicted_order_frame else [])
+            predicted_genus_frames.append(predicted_genus_frame if predicted_genus_frame else [])
 
             true_species_frame = []
             true_family_frame = []
-            true_order_frame = []
+            true_genus_frame = []
             
             if os.path.exists(label_path) and os.path.getsize(label_path) > 0:
                 with open(label_path, 'r') as f:
@@ -389,22 +392,22 @@ class TestTwoStage:
                         if species_idx < len(self.species_names):
                             species_name = self.species_names[species_idx]
                             
-                            if species_name in self.species_to_family:
-                                family_name = self.species_to_family[species_name]
-                                # Get the index of the family in the level_to_idx mapping
-                                if 2 in self.level_to_idx and family_name in self.level_to_idx[2]:
-                                    family_idx = self.level_to_idx[2][family_name]
-                                    true_family_frame.append([family_idx] + box_coords)
+                            if species_name in self.species_to_genus:
+                                genus_name = self.species_to_genus[species_name]
+                                # Get the index of the genus in the level_to_idx mapping
+                                if 2 in self.level_to_idx and genus_name in self.level_to_idx[2]:
+                                    genus_idx = self.level_to_idx[2][genus_name]
+                                    true_genus_frame.append([genus_idx] + box_coords)
                                 
-                                if family_name in self.family_to_order:
-                                    order_name = self.family_to_order[family_name]
-                                    if 1 in self.level_to_idx and order_name in self.level_to_idx[1]:
-                                        order_idx = self.level_to_idx[1][order_name]
-                                        true_order_frame.append([order_idx] + box_coords)
+                                if genus_name in self.genus_to_family:
+                                    family_name = self.genus_to_family[genus_name]
+                                    if 1 in self.level_to_idx and family_name in self.level_to_idx[1]:
+                                        family_idx = self.level_to_idx[1][family_name]
+                                        true_family_frame.append([family_idx] + box_coords)
 
             true_species_frames.append(true_species_frame if true_species_frame else [])
             true_family_frames.append(true_family_frame if true_family_frame else [])
-            true_order_frames.append(true_order_frame if true_order_frame else [])
+            true_genus_frames.append(true_genus_frame if true_genus_frame else [])
 
         end_time = time.time()  # End timing
         
@@ -416,42 +419,42 @@ class TestTwoStage:
             writer.writerow([
                 "Image Name", 
                 "True Species Detections", 
+                "True Genus Detections",
                 "True Family Detections",
-                "True Order Detections",
                 "Species Detections", 
-                "Family Detections", 
-                "Order Detections"
+                "Genus Detections", 
+                "Family Detections"
             ])
             
-            for image_name, true_species, true_family, true_order, species_pred, family_pred, order_pred in zip(
+            for image_name, true_species, true_genus, true_family, species_pred, genus_pred, family_pred in zip(
                 image_names, 
                 true_species_frames, 
+                true_genus_frames,
                 true_family_frames,
-                true_order_frames,
                 predicted_frames, 
-                predicted_family_frames, 
-                predicted_order_frames
+                predicted_genus_frames, 
+                predicted_family_frames
             ):
                 writer.writerow([
                     image_name, 
                     true_species, 
+                    true_genus,
                     true_family,
-                    true_order,
                     species_pred, 
-                    family_pred, 
-                    order_pred
+                    genus_pred, 
+                    family_pred
                 ])
         
         print(f"Results saved to {output_file}")
-        return predicted_frames, true_species_frames, end_time - start_time, predicted_family_frames, predicted_order_frames, true_family_frames, true_order_frames
+        return predicted_frames, true_species_frames, end_time - start_time, predicted_genus_frames, predicted_family_frames, true_genus_frames, true_family_frames
     
     def run(self, test_dir):
         results = self.get_frames(test_dir)
         predicted_frames, true_species_frames, total_time = results[0], results[1], results[2]
-        predicted_family_frames = results[3]
-        predicted_order_frames = results[4]
-        true_family_frames = results[5]
-        true_order_frames = results[6]
+        predicted_genus_frames = results[3]
+        predicted_family_frames = results[4]
+        true_genus_frames = results[5]
+        true_family_frames = results[6]
         
         num_frames = len(os.listdir(os.path.join(test_dir, 'images')))
         avg_time_per_frame = total_time / num_frames
@@ -461,29 +464,37 @@ class TestTwoStage:
         
         self.calculate_metrics(
             predicted_frames, true_species_frames, 
-            predicted_family_frames, true_family_frames,
-            predicted_order_frames, true_order_frames
+            predicted_genus_frames, true_genus_frames,
+            predicted_family_frames, true_family_frames
         )
 
     def calculate_metrics(self, predicted_species_frames, true_species_frames, 
-                         predicted_family_frames, true_family_frames, 
-                         predicted_order_frames, true_order_frames):
+                         predicted_genus_frames, true_genus_frames, 
+                         predicted_family_frames, true_family_frames):
         """Calculate metrics at all taxonomic levels"""
-        # Get list of species, families and orders
+        # Get list of species, families and genera using the same order as model training
         species_list = self.species_names
-        family_list = sorted(list(set(self.species_to_family.values())))
-        order_list = sorted(list(set(self.family_to_order.values())))
+        
+        # Use the index mappings from the model to ensure consistency
+        if 1 in self.idx_to_level and 2 in self.idx_to_level:
+            family_list = [self.idx_to_level[1][i] for i in sorted(self.idx_to_level[1].keys())]
+            genus_list = [self.idx_to_level[2][i] for i in sorted(self.idx_to_level[2].keys())]
+        else:
+            # Fallback to sorted lists (may cause issues)
+            print("Warning: Using fallback sorted lists for taxonomy - this may cause index mismatches")
+            genus_list = sorted(list(set(self.species_to_genus.values())))
+            family_list = sorted(list(set(self.genus_to_family.values())))
         
         # Print the index mappings we're using for evaluation
         print("\nUsing the following index mappings for evaluation:")
-        print("\nOrder indices:")
-        for i, order in enumerate(order_list):
-            print(f"  {i}: {order}")
-        
         print("\nFamily indices:")
         for i, family in enumerate(family_list):
             print(f"  {i}: {family}")
         
+        print("\nGenus indices:")
+        for i, genus in enumerate(genus_list):
+            print(f"  {i}: {genus}")
+        
         print("\nSpecies indices:")
         for i, species in enumerate(species_list):
             print(f"  {i}: {species}")
@@ -491,11 +502,11 @@ class TestTwoStage:
         # Dictionary to track prediction category counts for debugging
         prediction_counts = {
             "true_species_boxes": sum(len(frame) for frame in true_species_frames),
+            "true_genus_boxes": sum(len(frame) for frame in true_genus_frames),
             "true_family_boxes": sum(len(frame) for frame in true_family_frames),
-            "true_order_boxes": sum(len(frame) for frame in true_order_frames),
             "predicted_species": sum(len(frame) for frame in predicted_species_frames),
-            "predicted_family": sum(len(frame) for frame in predicted_family_frames),
-            "predicted_order": sum(len(frame) for frame in predicted_order_frames)
+            "predicted_genus": sum(len(frame) for frame in predicted_genus_frames),
+            "predicted_family": sum(len(frame) for frame in predicted_family_frames)
         }
         
         print(f"Prediction counts: {prediction_counts}")
@@ -504,11 +515,11 @@ class TestTwoStage:
         print("\n=== Species-level Metrics ===")
         self.get_metrics(predicted_species_frames, true_species_frames, species_list)
         
+        print("\n=== Genus-level Metrics ===")
+        self.get_metrics(predicted_genus_frames, true_genus_frames, genus_list)
+        
         print("\n=== Family-level Metrics ===")
         self.get_metrics(predicted_family_frames, true_family_frames, family_list)
-        
-        print("\n=== Order-level Metrics ===")
-        self.get_metrics(predicted_order_frames, true_order_frames, order_list)
 
     def get_metrics(self, predicted_frames, true_frames, labels):
         """Calculate metrics for object detection predictions"""
@@ -667,4 +678,4 @@ if __name__ == "__main__":
     hierarchical_model_path = "/mnt/nvme0n1p1/mit/two-stage-detection/hierarchical/hierarchical-weights.pth"
     output_directory = "./output"
     
-    test_multitask(species_names, test_directory, yolo_model_path, hierarchical_model_path, output_directory)
+    test(species_names, test_directory, yolo_model_path, hierarchical_model_path, output_directory)

bplusplus/tracker.py

@@ -0,0 +1,261 @@
+import numpy as np
+import uuid
+from scipy.optimize import linear_sum_assignment
+from collections import deque
+
+class BoundingBox:
+    def __init__(self, x, y, width, height, frame_id, track_id=None):
+        self.x = x
+        self.y = y
+        self.width = width
+        self.height = height
+        self.area = width * height
+        self.frame_id = frame_id
+        self.track_id = track_id
+    
+    def center(self):
+        return (self.x + self.width/2, self.y + self.height/2)
+    
+    @classmethod
+    def from_xyxy(cls, x1, y1, x2, y2, frame_id, track_id=None):
+        """Create BoundingBox from x1,y1,x2,y2 coordinates"""
+        width = x2 - x1
+        height = y2 - y1
+        return cls(x1, y1, width, height, frame_id, track_id)
+
+class InsectTracker:
+    def __init__(self, image_height, image_width, max_frames=30, w_dist=0.7, w_area=0.3, cost_threshold=0.8, track_memory_frames=None, debug=False):
+        self.image_height = image_height
+        self.image_width = image_width
+        self.max_dist = np.sqrt(image_height**2 + image_width**2)
+        self.max_frames = max_frames
+        self.w_dist = w_dist
+        self.w_area = w_area
+        self.cost_threshold = cost_threshold
+        self.debug = debug
+        
+        # If track_memory_frames not specified, use max_frames (full history window)
+        self.track_memory_frames = track_memory_frames if track_memory_frames is not None else max_frames
+        if self.debug:
+            print(f"DEBUG: Tracker initialized with max_frames={max_frames}, track_memory_frames={self.track_memory_frames}")
+        
+        self.tracking_history = deque(maxlen=max_frames)
+        self.current_tracks = []
+        self.lost_tracks = {}  # track_id -> {box: BoundingBox, frames_lost: int}
+    
+    def _generate_track_id(self):
+        """Generate a unique UUID for a new track"""
+        return str(uuid.uuid4())
+    
+    def calculate_cost(self, box1, box2):
+        """Calculate cost between two bounding boxes as per equation (4)"""
+        # Calculate center points
+        cx1, cy1 = box1.center()
+        cx2, cy2 = box2.center()
+        
+        # Euclidean distance (equation 1)
+        dist = np.sqrt((cx2 - cx1)**2 + (cy2 - cy1)**2)
+        
+        # Normalized distance (equation 2 used for normalization)
+        norm_dist = dist / self.max_dist
+        
+        # Area cost (equation 3)
+        min_area = min(box1.area, box2.area)
+        max_area = max(box1.area, box2.area)
+        area_cost = min_area / max_area if max_area > 0 else 1.0
+        
+        # Final cost (equation 4)
+        cost = (norm_dist * self.w_dist) + ((1 - area_cost) * self.w_area)
+        
+        return cost
+    
+    def build_cost_matrix(self, prev_boxes, curr_boxes):
+        """Build cost matrix for Hungarian algorithm"""
+        n_prev = len(prev_boxes)
+        n_curr = len(curr_boxes)
+        n = max(n_prev, n_curr)
+        
+        # Initialize cost matrix with high values
+        cost_matrix = np.ones((n, n)) * 999.0
+        
+        # Fill in actual costs
+        for i in range(n_prev):
+            for j in range(n_curr):
+                cost_matrix[i, j] = self.calculate_cost(prev_boxes[i], curr_boxes[j])
+        
+        return cost_matrix, n_prev, n_curr
+    
+    def update(self, new_detections, frame_id):
+        """
+        Update tracking with new detections from YOLO
+        
+        Args:
+            new_detections: List of YOLO detection boxes (x1, y1, x2, y2 format)
+            frame_id: Current frame number
+            
+        Returns:
+            List of track IDs corresponding to each detection
+        """
+        # Handle empty detection list (no detections in this frame)
+        if not new_detections:
+            if self.debug:
+                print(f"DEBUG: Frame {frame_id} has no detections")
+            # Move all current tracks to lost tracks
+            for track in self.current_tracks:
+                if track.track_id not in self.lost_tracks:
+                    self.lost_tracks[track.track_id] = {
+                        'box': track,
+                        'frames_lost': 1
+                    }
+                    if self.debug:
+                        print(f"DEBUG: Moved track {track.track_id} to lost tracks")
+                else:
+                    self.lost_tracks[track.track_id]['frames_lost'] += 1
+            
+            # Age lost tracks and remove old ones
+            self._age_lost_tracks()
+            
+            self.current_tracks = []
+            self.tracking_history.append([])
+            return []
+        
+        # Convert YOLO detections to BoundingBox objects
+        new_boxes = []
+        for i, detection in enumerate(new_detections):
+            x1, y1, x2, y2 = detection[:4]
+            bbox = BoundingBox.from_xyxy(x1, y1, x2, y2, frame_id)
+            new_boxes.append(bbox)
+        
+        # If this is the first frame or no existing tracks, assign new track IDs to all boxes
+        if not self.current_tracks and not self.lost_tracks:
+            track_ids = []
+            for box in new_boxes:
+                box.track_id = self._generate_track_id()
+                track_ids.append(box.track_id)
+                if self.debug:
+                    print(f"DEBUG: FIRST FRAME - Assigned track ID {box.track_id} to new detection")
+            self.current_tracks = new_boxes
+            self.tracking_history.append(new_boxes)
+            return track_ids
+        
+        # Combine current tracks and lost tracks for matching
+        all_previous_tracks = self.current_tracks.copy()
+        lost_track_list = []
+        
+        for track_id, lost_info in self.lost_tracks.items():
+            lost_track_list.append(lost_info['box'])
+            lost_track_list[-1].track_id = track_id  # Ensure track_id is preserved
+        
+        all_previous_tracks.extend(lost_track_list)
+        
+        if not all_previous_tracks:
+            # No previous tracks at all, assign new IDs
+            track_ids = []
+            for box in new_boxes:
+                box.track_id = self._generate_track_id()
+                track_ids.append(box.track_id)
+                if self.debug:
+                    print(f"DEBUG: No previous tracks - Assigned track ID {box.track_id} to new detection")
+            self.current_tracks = new_boxes
+            self.tracking_history.append(new_boxes)
+            return track_ids
+        
+        # Build cost matrix including lost tracks
+        cost_matrix, n_prev, n_curr = self.build_cost_matrix(all_previous_tracks, new_boxes)
+        
+        # Apply Hungarian algorithm
+        row_indices, col_indices = linear_sum_assignment(cost_matrix)
+        
+        # Assign track IDs based on the matching
+        assigned_curr_indices = set()
+        track_ids = [None] * len(new_boxes)
+        recovered_tracks = set()  # Track IDs that were recovered from lost tracks
+        
+        if self.debug:
+            print(f"DEBUG: Hungarian assignment - rows: {row_indices}, cols: {col_indices}")
+            print(f"DEBUG: Cost threshold: {self.cost_threshold}")
+            print(f"DEBUG: Current tracks: {len(self.current_tracks)}, Lost tracks: {len(self.lost_tracks)}")
+        
+        for i, j in zip(row_indices, col_indices):
+            # Only consider valid assignments (not dummy rows/columns)
+            if i < n_prev and j < n_curr:
+                cost = cost_matrix[i, j]
+                if self.debug:
+                    print(f"DEBUG: Checking assignment {i}->{j}, cost: {cost:.3f}")
+                # Check if cost is below threshold
+                if cost < self.cost_threshold:
+                    # Assign the track ID from previous box to current box
+                    prev_track_id = all_previous_tracks[i].track_id
+                    new_boxes[j].track_id = prev_track_id
+                    track_ids[j] = prev_track_id
+                    assigned_curr_indices.add(j)
+                    
+                    # Check if this was a lost track being recovered
+                    if prev_track_id in self.lost_tracks:
+                        recovered_tracks.add(prev_track_id)
+                        if self.debug:
+                            print(f"DEBUG: RECOVERED lost track ID {prev_track_id} for detection {j} (was lost for {self.lost_tracks[prev_track_id]['frames_lost']} frames)")
+                    else:
+                        if self.debug:
+                            print(f"DEBUG: Continued track ID {prev_track_id} for detection {j}")
+                else:
+                    if self.debug:
+                        print(f"DEBUG: Cost {cost:.3f} above threshold {self.cost_threshold}, not assigning")
+        
+        # Remove recovered tracks from lost tracks
+        for track_id in recovered_tracks:
+            del self.lost_tracks[track_id]
+        
+        # Assign new track IDs to unassigned current boxes (new insects)
+        for j in range(n_curr):
+            if j not in assigned_curr_indices:
+                new_boxes[j].track_id = self._generate_track_id()
+                track_ids[j] = new_boxes[j].track_id
+                if self.debug:
+                    print(f"DEBUG: Assigned NEW track ID {new_boxes[j].track_id} to detection {j}")
+        
+        # Move unmatched current tracks to lost tracks (tracks that disappeared this frame)
+        matched_track_ids = {track_ids[j] for j in assigned_curr_indices if track_ids[j] is not None}
+        for track in self.current_tracks:
+            if track.track_id not in matched_track_ids and track.track_id not in recovered_tracks:
+                if track.track_id not in self.lost_tracks:
+                    self.lost_tracks[track.track_id] = {
+                        'box': track,
+                        'frames_lost': 1
+                    }
+                    if self.debug:
+                        print(f"DEBUG: Track {track.track_id} disappeared, moved to lost tracks")
+        
+        # Age lost tracks and remove old ones
+        self._age_lost_tracks()
+        
+        # Update current tracks
+        self.current_tracks = new_boxes
+        
+        # Add to tracking history
+        self.tracking_history.append(new_boxes)
+        
+        return track_ids 
+    
+    def _age_lost_tracks(self):
+        """Age lost tracks and remove those that have been lost too long"""
+        tracks_to_remove = []
+        for track_id, lost_info in self.lost_tracks.items():
+            lost_info['frames_lost'] += 1
+            if lost_info['frames_lost'] > self.track_memory_frames:
+                tracks_to_remove.append(track_id)
+                if self.debug:
+                    print(f"DEBUG: Permanently removing track {track_id} (lost for {lost_info['frames_lost']} frames)")
+        
+        for track_id in tracks_to_remove:
+            del self.lost_tracks[track_id] 
+
+    def get_tracking_stats(self):
+        """Get current tracking statistics for debugging/monitoring"""
+        return {
+            'active_tracks': len(self.current_tracks),
+            'lost_tracks': len(self.lost_tracks),
+            'active_track_ids': [track.track_id for track in self.current_tracks],
+            'lost_track_ids': list(self.lost_tracks.keys()),
+            'total_history_frames': len(self.tracking_history)
+        }