Add template matching tracking to VideoEditor

This commit introduces a new tracking method using template matching, enhancing the tracking capabilities of the VideoEditor. It includes the ability to set a tracking template from a selected region, track the template in the current frame, and toggle template matching on and off. Additionally, debug messages have been added to provide insights during the template tracking process, improving user experience and functionality.
2025-09-26 14:24:31 +02:00
parent 71e5870306
commit c749d9af80
1 changed files with 146 additions and 1 deletions
--- a/croppa/main.py
+++ b/croppa/main.py
@@ -876,6 +876,13 @@ class VideoEditor:
        # Optical flow tracking
        self.optical_flow_enabled = False
        self.previous_frame_for_flow = None
+        
+        # Template matching tracking
+        self.template_matching_enabled = False
+        self.tracking_template = None
+        self.template_region = None  # (x, y, w, h) in rotated frame coordinates
+        self.template_selection_start = None
+        self.template_selection_rect = None

        # Project view mode
        self.project_view_mode = False
@@ -1620,7 +1627,32 @@ class VideoEditor:
    
    def _get_interpolated_tracking_position(self, frame_number):
        """Linear interpolation in ROTATED frame coords. Returns (rx, ry) or None."""
-        # First try feature tracking if enabled - but use smooth interpolation instead of averaging
+        # First try template matching if enabled (much better than optical flow)
+        if self.template_matching_enabled and self.tracking_template is not None:
+            if self.current_display_frame is not None:
+                # Apply transformations to get the display frame
+                display_frame = self.apply_crop_zoom_and_rotation(self.current_display_frame)
+                if display_frame is not None:
+                    # Track template in display frame
+                    result = self.track_template(display_frame)
+                    if result:
+                        center_x, center_y, confidence = result
+                        print(f"DEBUG: Template match found at ({center_x}, {center_y}) with confidence {confidence:.2f}")
+                        
+                        # Map from display frame coordinates to rotated frame coordinates
+                        frame_height, frame_width = display_frame.shape[:2]
+                        available_height = self.window_height - (0 if self.is_image_mode else self.TIMELINE_HEIGHT)
+                        start_y = (available_height - frame_height) // 2
+                        start_x = (self.window_width - frame_width) // 2
+                        
+                        screen_x = center_x + start_x
+                        screen_y = center_y + start_y
+                        
+                        # Map from screen coordinates to rotated frame coordinates
+                        rx, ry = self._map_screen_to_rotated(screen_x, screen_y)
+                        return (rx, ry)
+        
+        # Fall back to feature tracking if enabled - but use smooth interpolation instead of averaging
        if self.feature_tracker.tracking_enabled:
            # Get the nearest frames with features for smooth interpolation
            feature_frames = sorted(self.feature_tracker.features.keys())
@@ -2079,6 +2111,91 @@ class VideoEditor:
        interp_y = start_center[1] + alpha * (end_center[1] - start_center[1])
        
        return (interp_x, interp_y)
+    
+    def set_tracking_template(self, frame, region):
+        """Set a template region for tracking (much better than optical flow)"""
+        try:
+            x, y, w, h = region
+            self.tracking_template = frame[y:y+h, x:x+w].copy()
+            self.template_region = region
+            print(f"DEBUG: Set tracking template with region {region}")
+            return True
+        except Exception as e:
+            print(f"Error setting tracking template: {e}")
+            return False
+    
+    def track_template(self, frame):
+        """Track the template in the current frame"""
+        if self.tracking_template is None:
+            return None
+        
+        try:
+            # Convert to grayscale
+            if len(frame.shape) == 3:
+                gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+            else:
+                gray_frame = frame
+                
+            if len(self.tracking_template.shape) == 3:
+                gray_template = cv2.cvtColor(self.tracking_template, cv2.COLOR_BGR2GRAY)
+            else:
+                gray_template = self.tracking_template
+            
+            # Template matching
+            result = cv2.matchTemplate(gray_frame, gray_template, cv2.TM_CCOEFF_NORMED)
+            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
+            
+            # Only accept matches above threshold
+            if max_val > 0.6:  # Adjust threshold as needed
+                # Get template center
+                template_h, template_w = gray_template.shape
+                center_x = max_loc[0] + template_w // 2
+                center_y = max_loc[1] + template_h // 2
+                return (center_x, center_y, max_val)
+            else:
+                return None
+                
+        except Exception as e:
+            print(f"Error in template tracking: {e}")
+            return None
+    
+    def _set_template_from_region(self, screen_rect):
+        """Set template from selected region"""
+        x, y, w, h = screen_rect
+        print(f"DEBUG: Setting template from region ({x}, {y}, {w}, {h})")
+        
+        if self.current_display_frame is not None:
+            # Apply transformations to get the display frame
+            display_frame = self.apply_crop_zoom_and_rotation(self.current_display_frame)
+            if display_frame is not None:
+                # Map screen coordinates to display frame coordinates
+                frame_height, frame_width = display_frame.shape[:2]
+                available_height = self.window_height - (0 if self.is_image_mode else self.TIMELINE_HEIGHT)
+                start_y = (available_height - frame_height) // 2
+                start_x = (self.window_width - frame_width) // 2
+                
+                # Convert screen coordinates to display frame coordinates
+                display_x = x - start_x
+                display_y = y - start_y
+                display_w = w
+                display_h = h
+                
+                # Ensure region is within frame bounds
+                if (display_x >= 0 and display_y >= 0 and 
+                    display_x + display_w <= frame_width and 
+                    display_y + display_h <= frame_height):
+                    
+                    # Extract template from display frame
+                    template = display_frame[display_y:display_y+display_h, display_x:display_x+display_w]
+                    if template.size > 0:
+                        self.tracking_template = template.copy()
+                        self.template_region = (display_x, display_y, display_w, display_h)
+                        self.show_feedback_message(f"Template set from region ({display_w}x{display_h})")
+                        print(f"DEBUG: Template set with size {template.shape}")
+                    else:
+                        self.show_feedback_message("Template region too small")
+                else:
+                    self.show_feedback_message("Template region outside frame bounds")


    def apply_rotation(self, frame):
@@ -2938,6 +3055,26 @@ class VideoEditor:
                self.selective_feature_deletion_start = None
                self.selective_feature_deletion_rect = None
        
+        # Handle Alt+Right-click+drag for template region selection
+        if event == cv2.EVENT_RBUTTONDOWN and (flags & cv2.EVENT_FLAG_ALTKEY):
+            if not self.is_image_mode:
+                self.template_selection_start = (x, y)
+                self.template_selection_rect = None
+                print(f"DEBUG: Started template selection at ({x}, {y})")
+        
+        # Handle Alt+Right-click+drag for template region selection
+        if event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_ALTKEY) and self.template_selection_start:
+            if not self.is_image_mode:
+                start_x, start_y = self.template_selection_start
+                self.template_selection_rect = (min(start_x, x), min(start_y, y), abs(x - start_x), abs(y - start_y))
+        
+        # Handle Alt+Right-click release for template region selection
+        if event == cv2.EVENT_RBUTTONUP and (flags & cv2.EVENT_FLAG_ALTKEY) and self.template_selection_start:
+            if not self.is_image_mode and self.template_selection_rect:
+                self._set_template_from_region(self.template_selection_rect)
+                self.template_selection_start = None
+                self.template_selection_rect = None
+        
        # Handle right-click for tracking points (no modifiers)
        if event == cv2.EVENT_RBUTTONDOWN and not (flags & (cv2.EVENT_FLAG_CTRLKEY | cv2.EVENT_FLAG_SHIFTKEY)):
            if not self.is_image_mode:
@@ -3787,8 +3924,10 @@ class VideoEditor:
            print("  G: Clear all feature data")
            print("  H: Switch detector (SIFT/ORB)")
            print("  o: Toggle optical flow tracking")
+            print("  m: Toggle template matching tracking")
            print("  Shift+Right-click+drag: Extract features from selected region")
            print("  Ctrl+Right-click+drag: Delete features from selected region")
+            print("  Alt+Right-click+drag: Set template region for tracking")
            if len(self.video_files) > 1:
                print("  N: Next video")
                print("  n: Previous video")
@@ -4124,6 +4263,12 @@ class VideoEditor:
                
                self.show_feedback_message(f"Optical flow {'ON' if self.optical_flow_enabled else 'OFF'}")
                self.save_state()
+            elif key == ord("m"):
+                # Toggle template matching tracking
+                self.template_matching_enabled = not self.template_matching_enabled
+                print(f"DEBUG: Template matching toggled to {self.template_matching_enabled}")
+                self.show_feedback_message(f"Template matching {'ON' if self.template_matching_enabled else 'OFF'}")
+                self.save_state()
            elif key == ord("t"):
                # Marker looping only for videos
                if not self.is_image_mode: