Update reid.py

reid.py

@@ -1,5 +1,5 @@
 """
-Enhanced ReID with MegaDescriptor-B-224 and adaptive strategies
+Simplified ReID with MegaDescriptor-L-384 (Largest Model)
 """
 import numpy as np
 import cv2
@@ -20,89 +20,63 @@ class DogFeatures:
     bbox: List[float] = field(default_factory=list)
     confidence: float = 0.5
     frame_num: int = 0
-    track_id: int = 0  # Add track ID for continuity
 
 
 class MegaDescriptorReID:
     """
-    Enhanced ReID with adaptive strategies and debugging
+    Simplified ReID using MegaDescriptor-L-384
     """
     
-    def __init__(self, device: str = 'cuda', max_expected_dogs: int = 10):
+    def __init__(self, device: str = 'cuda'):
         self.device = device if torch.cuda.is_available() else 'cpu'
-        self.base_threshold = 0.40
-        self.max_expected_dogs = max_expected_dogs
+        self.base_threshold = 0.35  # Lower default for L model
         
         # Dog database (temporary only)
         self.dog_database = {}  # dog_id -> list of DogFeatures
         self.next_dog_id = 1
         self.current_frame = 0
         
-        # Track continuity mapping
-        self.track_to_dog = {}  # track_id -> dog_id
-        self.dog_last_track = {}  # dog_id -> last_track_id
-        
         # Statistics for debugging
         self.match_stats = {
-            'new_dogs_created': 0,
-            'successful_matches': 0,
-            'threshold_adjustments': 0,
-            'track_continuity_matches': 0,
-            'similarity_scores': []
+            'new_dogs': [],
+            'matches': [],
+            'all_scores': []
         }
         
-        # Initialize MegaDescriptor
+        # Initialize MegaDescriptor-L
         self._initialize_megadescriptor()
         
-        print(f"✅ Enhanced MegaDescriptor ReID initialized")
-        print(f"   Device: {self.device}")
-        print(f"   Max expected dogs: {self.max_expected_dogs}")
+        print(f"✅ MegaDescriptor-L-384 ReID initialized on {self.device}")
         
     def _initialize_megadescriptor(self):
-        """Initialize MegaDescriptor-B-224"""
+        """Initialize MegaDescriptor-L-384 (Largest model)"""
         try:
+            # Load the largest MegaDescriptor model
+            print("📥 Loading MegaDescriptor-L-384 (this may take a moment)...")
             self.model = timm.create_model(
-                'hf-hub:BVRA/MegaDescriptor-B-224',
+                'hf-hub:BVRA/MegaDescriptor-L-384',
                 pretrained=True
             )
             self.model.to(self.device).eval()
             
-            # Get the preprocessing config
-            data_config = timm.data.resolve_model_data_config(self.model)
-            self.transform = timm.data.create_transform(**data_config, is_training=False)
+            # L model uses 384x384 input
+            self.transform = timm.data.create_transform(
+                input_size=(384, 384),
+                is_training=False,
+                mean=[0.5, 0.5, 0.5],
+                std=[0.5, 0.5, 0.5]
+            )
             
-            print("✅ MegaDescriptor-B-224 loaded successfully")
+            print("✅ MegaDescriptor-L-384 loaded successfully")
+            print("   • Model: Large (384x384 input)")
+            print("   • Features: 1024-dim")
             
         except Exception as e:
             print(f"❌ MegaDescriptor initialization error: {e}")
             self.model = None
     
-    def get_adaptive_threshold(self):
-        """Calculate adaptive threshold based on current dog count"""
-        num_dogs = len(self.dog_database)
-        original_threshold = self.base_threshold
-        
-        if num_dogs >= self.max_expected_dogs * 1.5:
-            # Way too many dogs, be very lenient
-            adapted = 0.25
-        elif num_dogs >= self.max_expected_dogs:
-            # Too many dogs, lower threshold
-            adapted = 0.30
-        elif num_dogs >= self.max_expected_dogs * 0.8:
-            # Approaching limit, start lowering
-            adapted = 0.35
-        else:
-            # Normal range
-            adapted = self.base_threshold
-        
-        if adapted != original_threshold:
-            self.match_stats['threshold_adjustments'] += 1
-            print(f"📊 Adaptive threshold: {original_threshold:.2f} → {adapted:.2f} (dogs: {num_dogs})")
-        
-        return adapted
-    
-    def extract_features(self, image: np.ndarray, bbox: List[float] = None, track_id: int = None) -> Optional[DogFeatures]:
-        """Extract features using MegaDescriptor"""
+    def extract_features(self, image: np.ndarray, bbox: List[float] = None) -> Optional[DogFeatures]:
+        """Extract features using MegaDescriptor-L"""
         if image is None or image.size == 0 or self.model is None:
             return None
         
@@ -114,7 +88,7 @@ class MegaDescriptorReID:
             from PIL import Image
             pil_img = Image.fromarray(img_rgb)
             
-            # Apply MegaDescriptor transforms
+            # Apply transforms (384x384 for L model)
             img_tensor = self.transform(pil_img).unsqueeze(0).to(self.device)
             
             # Extract features
@@ -127,8 +101,7 @@ class MegaDescriptorReID:
             return DogFeatures(
                 features=features,
                 bbox=bbox if bbox else [0, 0, 100, 100],
-                frame_num=self.current_frame,
-                track_id=track_id if track_id else 0
+                frame_num=self.current_frame
             )
             
         except Exception as e:
@@ -136,28 +109,9 @@ class MegaDescriptorReID:
             return None
     
     def match_or_register(self, track, image_crop=None) -> Tuple[int, float]:
-        """Match or register a dog with enhanced strategies"""
+        """Simple match or register without complex strategies"""
         self.current_frame += 1
         
-        # Get track ID for continuity
-        track_id = track.track_id if hasattr(track, 'track_id') else 0
-        
-        # Check if this track already has a dog ID (continuity)
-        if track_id in self.track_to_dog:
-            dog_id = self.track_to_dog[track_id]
-            self.match_stats['track_continuity_matches'] += 1
-            print(f"  🔄 Track continuity: Track {track_id} → Dog {dog_id}")
-            
-            # Still extract and store features for future matching
-            for det in reversed(track.detections[-3:]):
-                if det.image_crop is not None:
-                    features = self.extract_features(det.image_crop, det.bbox, track_id)
-                    if features:
-                        self._update_dog_features(dog_id, features)
-                    break
-            
-            return dog_id, 0.95  # High confidence for track continuity
-        
         # Get detection with image
         detection = None
         for det in reversed(track.detections[-3:]):
@@ -172,8 +126,7 @@ class MegaDescriptorReID:
         # Extract features
         features = self.extract_features(
             image_crop, 
-            detection.bbox if hasattr(detection, 'bbox') else None,
-            track_id
+            detection.bbox if hasattr(detection, 'bbox') else None
         )
         
         if features is None:
@@ -181,18 +134,15 @@ class MegaDescriptorReID:
         
         features.confidence = detection.confidence if hasattr(detection, 'confidence') else 0.5
         
-        # Get adaptive threshold
-        threshold = self.get_adaptive_threshold()
-        
         # Find best match
         best_dog_id = None
         best_score = -1.0
-        all_scores = []  # For debugging
+        debug_scores = []
         
         for dog_id, dog_features_list in self.dog_database.items():
             # Calculate similarity with stored features
             similarities = []
-            for stored_feat in dog_features_list[-10:]:  # Use last 10 features
+            for stored_feat in dog_features_list[-20:]:  # Use last 20 features
                 sim = cosine_similarity(
                     features.features.reshape(1, -1),
                     stored_feat.features.reshape(1, -1)
@@ -200,104 +150,53 @@ class MegaDescriptorReID:
                 similarities.append(sim)
             
             if similarities:
-                avg_similarity = np.mean(similarities)
-                max_similarity = np.max(similarities)
-                # Use weighted combination
-                weighted_sim = 0.7 * avg_similarity + 0.3 * max_similarity
+                # Use max similarity for L model (more discriminative)
+                max_sim = np.max(similarities)
+                avg_sim = np.mean(similarities)
+                # Weight max more for L model
+                final_score = 0.6 * max_sim + 0.4 * avg_sim
                 
-                all_scores.append((dog_id, weighted_sim))
+                debug_scores.append((dog_id, final_score, max_sim, avg_sim))
                 
-                if weighted_sim > best_score:
-                    best_score = weighted_sim
+                if final_score > best_score:
+                    best_score = final_score
                     best_dog_id = dog_id
         
         # Debug output
-        if all_scores:
-            self.match_stats['similarity_scores'].extend([s[1] for s in all_scores])
-            top_matches = sorted(all_scores, key=lambda x: x[1], reverse=True)[:3]
-            print(f"  🔍 Top matches: {[(f'Dog{d}', f'{s:.3f}') for d, s in top_matches]}")
+        if debug_scores:
+            self.match_stats['all_scores'].append(best_score)
+            top_matches = sorted(debug_scores, key=lambda x: x[1], reverse=True)[:3]
+            print(f"  🔍 Frame {self.current_frame} matches:")
+            for dog_id, final, max_s, avg_s in top_matches[:3]:
+                print(f"     Dog {dog_id}: {final:.3f} (max:{max_s:.3f}, avg:{avg_s:.3f})")
+        
+        # Simple decision with threshold
+        threshold = self.base_threshold
+        print(f"  📊 Best score: {best_score:.3f}, Threshold: {threshold:.3f}")
         
-        # Decision: match or new dog
         if best_dog_id is not None and best_score >= threshold:
             # Match found
-            self.match_stats['successful_matches'] += 1
-            print(f"  ✅ Matched to Dog {best_dog_id} (score: {best_score:.3f}, threshold: {threshold:.3f})")
-            
-            # Update track mapping
-            self.track_to_dog[track_id] = best_dog_id
-            self.dog_last_track[best_dog_id] = track_id
+            self.match_stats['matches'].append((best_dog_id, best_score))
+            print(f"  ✅ Matched to Dog {best_dog_id}")
             
             # Update features
-            self._update_dog_features(best_dog_id, features)
+            self.dog_database[best_dog_id].append(features)
+            # Keep last 30 features
+            if len(self.dog_database[best_dog_id]) > 30:
+                self.dog_database[best_dog_id] = self.dog_database[best_dog_id][-30:]
             
             return best_dog_id, best_score
         else:
-            # Check if we should be more aggressive due to dog count
-            if len(self.dog_database) >= self.max_expected_dogs and best_score > 0.2:
-                # Force match if we have too many dogs and score is reasonable
-                print(f"  ⚠️ Forced match to Dog {best_dog_id} (too many dogs, score: {best_score:.3f})")
-                self.track_to_dog[track_id] = best_dog_id
-                self._update_dog_features(best_dog_id, features)
-                return best_dog_id, best_score
-            
             # New dog
-            new_dog_id = self._register_new_dog(features, track_id)
+            new_dog_id = self.next_dog_id
+            self.next_dog_id += 1
+            self.match_stats['new_dogs'].append(new_dog_id)
+            
+            self.dog_database[new_dog_id] = [features]
+            print(f"  🆕 New dog: Dog {new_dog_id} (Total: {len(self.dog_database)})")
+            
             return new_dog_id, 1.0
     
-    def _update_dog_features(self, dog_id: int, features: DogFeatures):
-        """Update dog features database"""
-        self.dog_database[dog_id].append(features)
-        # Keep more features for better matching
-        if len(self.dog_database[dog_id]) > 30:
-            self.dog_database[dog_id] = self.dog_database[dog_id][-30:]
-    
-    def _register_new_dog(self, features: DogFeatures, track_id: int) -> int:
-        """Register a new dog"""
-        new_dog_id = self.next_dog_id
-        self.next_dog_id += 1
-        self.match_stats['new_dogs_created'] += 1
-        
-        self.dog_database[new_dog_id] = [features]
-        self.track_to_dog[track_id] = new_dog_id
-        self.dog_last_track[new_dog_id] = track_id
-        
-        print(f"  🆕 New dog registered: Dog {new_dog_id} (Total: {len(self.dog_database)})")
-        
-        return new_dog_id
-    
-    def post_process_merge(self, merge_threshold: float = 0.7):
-        """Post-process to merge similar dogs"""
-        print("\n🔄 Post-processing: Checking for similar dogs to merge...")
-        
-        merged_count = 0
-        dog_ids = list(self.dog_database.keys())
-        
-        for i, dog1_id in enumerate(dog_ids):
-            if dog1_id not in self.dog_database:
-                continue
-                
-            for dog2_id in dog_ids[i+1:]:
-                if dog2_id not in self.dog_database:
-                    continue
-                
-                # Compare average features
-                feat1 = np.mean([f.features for f in self.dog_database[dog1_id]], axis=0)
-                feat2 = np.mean([f.features for f in self.dog_database[dog2_id]], axis=0)
-                
-                similarity = cosine_similarity(feat1.reshape(1, -1), feat2.reshape(1, -1))[0, 0]
-                
-                if similarity > merge_threshold:
-                    # Merge dog2 into dog1
-                    print(f"  🔗 Merging Dog {dog2_id} into Dog {dog1_id} (similarity: {similarity:.3f})")
-                    self.dog_database[dog1_id].extend(self.dog_database[dog2_id])
-                    del self.dog_database[dog2_id]
-                    merged_count += 1
-        
-        if merged_count > 0:
-            print(f"  ✅ Merged {merged_count} duplicate dogs. Final count: {len(self.dog_database)}")
-        
-        return merged_count
-    
     def match_or_register_all(self, track) -> Dict:
         """Compatible interface"""
         dog_id, confidence = self.match_or_register(track)
@@ -309,36 +208,39 @@ class MegaDescriptorReID:
         }
     
     def set_all_thresholds(self, threshold: float):
-        """Update base threshold"""
+        """Update threshold"""
         self.base_threshold = max(0.15, min(0.95, threshold))
-        print(f"📊 Base ReID threshold set to: {self.base_threshold:.2f}")
+        print(f"📊 ReID threshold set to: {self.base_threshold:.2f}")
     
     def reset_all(self):
         """Reset for new video"""
+        # Print final statistics
+        if self.dog_database:
+            print("\n" + "="*50)
+            print("📈 Final Session Statistics:")
+            print(f"  • Total dogs detected: {len(self.dog_database)}")
+            print(f"  • New dog creations: {len(self.match_stats['new_dogs'])}")
+            print(f"  • Successful matches: {len(self.match_stats['matches'])}")
+            
+            if self.match_stats['all_scores']:
+                scores = self.match_stats['all_scores']
+                print(f"  • Match scores - Avg: {np.mean(scores):.3f}, Min: {np.min(scores):.3f}, Max: {np.max(scores):.3f}")
+            
+            print("\n  Dogs summary:")
+            for dog_id, features_list in self.dog_database.items():
+                print(f"    Dog {dog_id}: {len(features_list)} features stored")
+            print("="*50 + "\n")
+        
+        # Clear everything
         self.dog_database.clear()
-        self.track_to_dog.clear()
-        self.dog_last_track.clear()
         self.next_dog_id = 1
         self.current_frame = 0
         
-        # Print debug statistics before reset
-        if self.match_stats['new_dogs_created'] > 0:
-            print("\n📈 Session Statistics:")
-            print(f"  • New dogs created: {self.match_stats['new_dogs_created']}")
-            print(f"  • Successful matches: {self.match_stats['successful_matches']}")
-            print(f"  • Track continuity matches: {self.match_stats['track_continuity_matches']}")
-            print(f"  • Threshold adjustments: {self.match_stats['threshold_adjustments']}")
-            if self.match_stats['similarity_scores']:
-                scores = self.match_stats['similarity_scores']
-                print(f"  • Avg similarity: {np.mean(scores):.3f} (min: {np.min(scores):.3f}, max: {np.max(scores):.3f})")
-        
         # Reset statistics
         self.match_stats = {
-            'new_dogs_created': 0,
-            'successful_matches': 0,
-            'threshold_adjustments': 0,
-            'track_continuity_matches': 0,
-            'similarity_scores': []
+            'new_dogs': [],
+            'matches': [],
+            'all_scores': []
         }
         
         print("🔄 ReID reset\n")
@@ -348,17 +250,8 @@ class MegaDescriptorReID:
         return {
             'total_dogs': len(self.dog_database),
             'threshold': self.base_threshold,
-            'stats': self.match_stats
+            'model': 'MegaDescriptor-L-384'
         }
-    def aggressive_merge(self):
-        """Keep merging until no more merges possible"""
-        total_merged = 0
-        while True:
-            merged = self.post_process_merge(merge_threshold=0.5)
-            if merged == 0:
-                break
-            total_merged += merged
-        return total_merged
 
 
 # Compatibility aliases