bplusplus 1.2.2__py3-none-any.whl → 1.2.4__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.
     
@@ -74,6 +74,16 @@ def setup_gpu():
         logger.warning("Falling back to CPU")
         return torch.device("cpu")
 
+# Add this check for backwards compatibility
+if hasattr(torch.serialization, 'add_safe_globals'):
+    torch.serialization.add_safe_globals([
+        'torch.LongTensor',
+        'torch.cuda.LongTensor',
+        'torch.FloatStorage',
+        'torch.FloatStorage',
+        'torch.cuda.FloatStorage',
+    ])
+
 class HierarchicalInsectClassifier(nn.Module):
     def __init__(self, num_classes_per_level):
         """
@@ -243,8 +253,15 @@ 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_genus, genus_to_family = 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_genus, genus_to_family = get_taxonomy(species_names)
         else:
@@ -285,8 +302,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])
         ])
@@ -467,10 +482,18 @@ class TestTwoStage:
                          predicted_genus_frames, true_genus_frames, 
                          predicted_family_frames, true_family_frames):
         """Calculate metrics at all taxonomic levels"""
-        # Get list of species, families and genera
+        # Get list of species, families and genera using the same order as model training
         species_list = self.species_names
-        genus_list = sorted(list(set(self.species_to_genus.values())))
-        family_list = sorted(list(set(self.genus_to_family.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:")
@@ -665,4 +688,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)
+        } 

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

@@ -14,18 +14,28 @@ import logging
 from tqdm import tqdm  
 import sys
 
-def train_multitask(batch_size=4, epochs=30, patience=3, img_size=640, data_dir='/mnt/nvme0n1p1/datasets/insect/bjerge-train2', output_dir='./output', species_list=None):
+def train(batch_size=4, epochs=30, patience=3, img_size=640, data_dir='input', output_dir='./output', species_list=None, num_workers=4):
     """
     Main function to run the entire training pipeline.
     Sets up datasets, model, training process and handles errors.
+    
+    Args:
+        batch_size (int): Number of samples per batch. Default: 4
+        epochs (int): Maximum number of training epochs. Default: 30
+        patience (int): Early stopping patience (epochs without improvement). Default: 3
+        img_size (int): Target image size for training. Default: 640
+        data_dir (str): Directory containing train/valid subdirectories. Default: 'input'
+        output_dir (str): Directory to save trained model and logs. Default: './output'
+        species_list (list): List of species names for training. Required.
+        num_workers (int): Number of DataLoader worker processes. 
+                          Set to 0 to disable multiprocessing (most stable). Default: 4
     """
     global logger, device
 
     logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
     logger = logging.getLogger(__name__)
 
-    logger.info(f"Hyperparameters - Batch size: {batch_size}, Epochs: {epochs}, Patience: {patience}, Image size: {img_size}, Data directory: {data_dir}, Output directory: {output_dir}")
-
+    logger.info(f"Hyperparameters - Batch size: {batch_size}, Epochs: {epochs}, Patience: {patience}, Image size: {img_size}, Data directory: {data_dir}, Output directory: {output_dir}, Num workers: {num_workers}")
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -52,7 +62,7 @@ def train_multitask(batch_size=4, epochs=30, patience=3, img_size=640, data_dir=
     
     taxonomy = get_taxonomy(species_list)
     
-    level_to_idx, parent_child_relationship = create_mappings(taxonomy)
+    level_to_idx, parent_child_relationship = create_mappings(taxonomy, species_list)
     
     num_classes_per_level = [len(taxonomy[level]) if isinstance(taxonomy[level], list) 
                             else len(taxonomy[level].keys()) for level in sorted(taxonomy.keys())]
@@ -75,14 +85,14 @@ def train_multitask(batch_size=4, epochs=30, patience=3, img_size=640, data_dir=
         train_dataset,
         batch_size=batch_size,
         shuffle=True,
-        num_workers=4
+        num_workers=num_workers
     )
     
     val_loader = DataLoader(
         val_dataset,
         batch_size=batch_size,
         shuffle=False,
-        num_workers=4
+        num_workers=num_workers
     )
     
     try:
@@ -150,14 +160,17 @@ def get_taxonomy(species_list):
     species_to_genus = {}
     genus_to_family = {}
     
-    logger.info(f"Building taxonomy from GBIF for {len(species_list)} species")
+    species_list_for_gbif = [s for s in species_list if s.lower() != 'unknown']
+    has_unknown = len(species_list_for_gbif) != len(species_list)
+    
+    logger.info(f"Building taxonomy from GBIF for {len(species_list_for_gbif)} species")
     
     print("\nTaxonomy Results:")
     print("-" * 80)
     print(f"{'Species':<30} {'Family':<20} {'Genus':<20} {'Status'}")
     print("-" * 80)
     
-    for species_name in species_list:
+    for species_name in species_list_for_gbif:
         url = f"https://api.gbif.org/v1/species/match?name={species_name}&verbose=true"
         try:
             response = requests.get(url)
@@ -199,6 +212,19 @@ def get_taxonomy(species_list):
             print(f"{species_name:<30} {'Error':<20} {'Error':<20} FAILED")
             print(f"Error: {error_msg}")
             sys.exit(1)  # Stop the script
+
+    if has_unknown:
+        unknown_family = "Unknown"
+        unknown_genus = "Unknown"
+        unknown_species = "unknown"
+        
+        if unknown_family not in taxonomy[1]:
+            taxonomy[1].append(unknown_family)
+        
+        taxonomy[2][unknown_genus] = unknown_family
+        taxonomy[3][unknown_species] = unknown_genus
+        
+        print(f"{unknown_species:<30} {unknown_family:<20} {unknown_genus:<20} {'OK'}")
     
     taxonomy[1] = sorted(list(set(taxonomy[1])))
     print("-" * 80)
@@ -212,7 +238,7 @@ def get_taxonomy(species_list):
         print(f"  {i}: {family}")
     
     print("\nGenus indices:")
-    for i, genus in enumerate(taxonomy[2].keys()):
+    for i, genus in enumerate(sorted(taxonomy[2].keys())):
         print(f"  {i}: {genus}")
     
     print("\nSpecies indices:")
@@ -244,7 +270,7 @@ def get_species_from_directory(train_dir):
     logger.info(f"Found {len(species_list)} species in {train_dir}")
     return species_list
 
-def create_mappings(taxonomy):
+def create_mappings(taxonomy, species_list=None):
     """
     Creates mapping dictionaries from taxonomy data.
     Returns level-to-index mapping and parent-child relationships between taxonomic levels.
@@ -254,9 +280,17 @@ def create_mappings(taxonomy):
 
     for level, labels in taxonomy.items():
         if isinstance(labels, list):
+            # Level 1: Family (already sorted)
             level_to_idx[level] = {label: idx for idx, label in enumerate(labels)}
-        else:
-            level_to_idx[level] = {label: idx for idx, label in enumerate(labels.keys())}
+        else:  # dict for levels 2 and 3
+            if level == 3 and species_list is not None:
+                # For species, the order is determined by species_list
+                level_to_idx[level] = {label: idx for idx, label in enumerate(species_list)}
+            else:
+                # For genus (and as a fallback for species), sort alphabetically
+                sorted_keys = sorted(labels.keys())
+                level_to_idx[level] = {label: idx for idx, label in enumerate(sorted_keys)}
+            
             for child, parent in labels.items():
                 if (level, parent) not in parent_child_relationship:
                     parent_child_relationship[(level, parent)] = []
@@ -670,7 +704,7 @@ if __name__ == '__main__':
     species_list = [
         "Coccinella septempunctata", "Apis mellifera", "Bombus lapidarius", "Bombus terrestris",
         "Eupeodes corollae", "Episyrphus balteatus", "Aglais urticae", "Vespula vulgaris",
-        "Eristalis tenax"
+        "Eristalis tenax", "unknown"
     ]
-    train_multitask(species_list=species_list, epochs=2)
+    train(species_list=species_list, epochs=2)