Update app.py

app.py

@@ -63,7 +63,7 @@ class DatasetCollectionApp:
         print(f"Database has {len(self.db.get_all_dogs())} dogs")
 
     def create_visualization_video(self, video_path: str, sample_rate: int) -> str:
-        """Create visualization video with colored tracking boxes and IDs"""
+        """Create visualization video with boxes matching ReID temp_ids with ID labels"""
         try:
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
@@ -76,7 +76,7 @@ class DatasetCollectionApp:
             
             output_path = f"visualization_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
             
-            # Try multiple codecs in order of preference
+            # Try multiple codecs
             codecs = ['mp4v', 'XVID', 'MJPG']
             out = None
             
@@ -91,23 +91,30 @@ class DatasetCollectionApp:
                     out.release()
                     out = None
             
-            # If no codec worked, return None
             if out is None or not out.isOpened():
                 print("ERROR: Could not initialize VideoWriter with any codec")
                 cap.release()
                 return None
             
+            # Create separate tracker for visualization
             viz_tracker = DeepSORTTracker(
                 max_iou_distance=0.5,
                 max_age=90,
                 n_init=1,
-                use_appearance=False
+                use_appearance=True
             )
             
+            # Create separate ReID instance for visualization
+            viz_reid = SimplifiedReID(device=self.device)
+            viz_reid.set_threshold(self.reid.threshold)
+            viz_reid.set_video_source(video_path)
+            
+            # Color mapping for temp_ids
             track_colors = {}
+            
             frame_num = 0
             
-            print("\nCreating visualization video...")
+            print("\nCreating visualization video with ReID matching and ID labels...")
             
             while cap.isOpened():
                 ret, frame = cap.read()
@@ -115,47 +122,60 @@ class DatasetCollectionApp:
                     break
                 
                 if frame_num % sample_rate == 0:
+                    # Detect dogs
                     detections = self.detector.detect(frame)
+                    
+                    # Track dogs
                     tracks = viz_tracker.update(detections)
                     
+                    # Match with ReID to get temp_ids
                     for track in tracks:
-                        if track.track_id not in track_colors:
-                            track_colors[track.track_id] = (
+                        # Get ReID temp_id (matches gallery)
+                        result = viz_reid.match_or_register(track)
+                        temp_id = result['temp_id']
+                        
+                        if temp_id == 0:
+                            continue
+                        
+                        # Assign color based on temp_id
+                        if temp_id not in track_colors:
+                            track_colors[temp_id] = (
                                 np.random.randint(50, 255),
                                 np.random.randint(50, 255),
                                 np.random.randint(50, 255)
                             )
                         
                         x1, y1, x2, y2 = map(int, track.bbox)
-                        color = track_colors[track.track_id]
+                        color = track_colors[temp_id]
                         
-                        # Bold box (thickness = 6)
+                        # Draw bold box
                         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 6)
                         
                         # Add ID label
-                        label = f"ID: {track.track_id}"
+                        label = f"#{temp_id}"
                         font = cv2.FONT_HERSHEY_SIMPLEX
-                        font_scale = 1.2
+                        font_scale = 1.5
                         font_thickness = 3
                         
+                        # Get text size
                         (text_width, text_height), baseline = cv2.getTextSize(
                             label, font, font_scale, font_thickness
                         )
                         
-                        # Background rectangle for text
+                        # Draw filled rectangle background for text
                         cv2.rectangle(
                             frame,
-                            (x1, y1 - text_height - 10),
-                            (x1 + text_width + 10, y1),
+                            (x1, y1 - text_height - 15),
+                            (x1 + text_width + 15, y1),
                             color,
                             -1
                         )
                         
-                        # White text
+                        # Draw white text
                         cv2.putText(
                             frame,
                             label,
-                            (x1 + 5, y1 - 5),
+                            (x1 + 7, y1 - 7),
                             font,
                             font_scale,
                             (255, 255, 255),
@@ -172,9 +192,10 @@ class DatasetCollectionApp:
             cap.release()
             out.release()
             
-            # Verify file was created and has content
+            # Verify file
             if os.path.exists(output_path) and os.path.getsize(output_path) > 1000:
                 print(f"Visualization video saved: {output_path}")
+                print(f"Gallery temp_ids shown: {sorted(track_colors.keys())}")
                 return output_path
             else:
                 print("ERROR: Video file not created or is empty")