From 840440eb1acdbc3a563a32ee75119738dcc8838d Mon Sep 17 00:00:00 2001 From: PhatPhuckDave Date: Fri, 26 Sep 2025 17:23:04 +0200 Subject: [PATCH] Add multiple templates management to VideoEditor This commit introduces a system for managing multiple templates within the VideoEditor. Users can now add, remove, and navigate through templates, enhancing the flexibility of template tracking. The state management has been updated to save and load multiple templates, including their regions and frame ranges. Additionally, visual indicators for active templates have been implemented in the frame display, improving user feedback during editing sessions. --- croppa/main.py | 231 ++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 221 insertions(+), 10 deletions(-) diff --git a/croppa/main.py b/croppa/main.py index 8a7d97d..7f3a08f 100644 --- a/croppa/main.py +++ b/croppa/main.py @@ -885,6 +885,11 @@ class VideoEditor: self.multi_scale_template_matching = False # Disable multi-scale by default # (x, y, w, h) in rotated frame coordinates self.template_selection_start = None self.template_selection_rect = None + + # Multiple templates system + self.templates = {} # {template_id: {'template': image, 'region': (x,y,w,h), 'start_frame': int, 'end_frame': int}} + self.current_template_id = None + self.template_id_counter = 0 # Project view mode self.project_view_mode = False @@ -934,7 +939,15 @@ class VideoEditor: 'feature_tracker': self.feature_tracker.get_state_dict(), 'template_matching_enabled': self.template_matching_enabled, 'template_region': self.template_region, - 'multi_scale_template_matching': self.multi_scale_template_matching + 'multi_scale_template_matching': self.multi_scale_template_matching, + 'templates': {str(k): { + 'template': None, # Don't save template images (too large) + 'region': v['region'], + 'start_frame': v['start_frame'], + 'end_frame': v['end_frame'] + } for k, v in self.templates.items()}, + 'current_template_id': self.current_template_id, + 'template_id_counter': self.template_id_counter } with open(state_file, 'w') as f: @@ -1031,6 +1044,23 @@ class VideoEditor: self.tracking_template = None if 'multi_scale_template_matching' in state: self.multi_scale_template_matching = state['multi_scale_template_matching'] # Will be recreated on first use + + # Load multiple templates state + if 'templates' in state: + self.templates = {} + for template_id_str, template_data in state['templates'].items(): + template_id = int(template_id_str) + self.templates[template_id] = { + 'template': None, # Will be recreated when needed + 'region': template_data['region'], + 'start_frame': template_data['start_frame'], + 'end_frame': template_data['end_frame'] + } + print(f"Loaded {len(self.templates)} templates") + if 'current_template_id' in state: + self.current_template_id = state['current_template_id'] + if 'template_id_counter' in state: + self.template_id_counter = state['template_id_counter'] # Validate cut markers against current video length if self.cut_start_frame is not None and self.cut_start_frame >= self.total_frames: @@ -2222,12 +2252,12 @@ class VideoEditor: def track_template(self, frame): """Track the template in the current frame""" + # First, try to select the best template for current frame + if not self._select_best_template_for_frame(self.current_frame): + return None + if self.tracking_template is None: - # Try to recreate template from saved region - if self.template_region is not None: - self._recreate_template_from_region(frame) - if self.tracking_template is None: - return None + return None try: # Apply image preprocessing for better template matching @@ -2383,15 +2413,142 @@ class VideoEditor: # Extract template from raw frame template = self.current_display_frame[raw_y:raw_y+raw_h, raw_x:raw_x+raw_w] if template.size > 0: - self.tracking_template = template.copy() - self.template_region = (raw_x, raw_y, raw_w, raw_h) - self.show_feedback_message(f"Template set from region ({raw_w}x{raw_h})") - print(f"DEBUG: Template set with size {template.shape}") + # Add template to collection instead of setting single template + template_id = self.add_template(template, (raw_x, raw_y, raw_w, raw_h)) + self.show_feedback_message(f"Template {template_id} set from region ({raw_w}x{raw_h})") + print(f"DEBUG: Template {template_id} 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 add_template(self, template, region, start_frame=None, end_frame=None): + """Add a new template to the collection""" + template_id = self.template_id_counter + self.template_id_counter += 1 + + if start_frame is None: + start_frame = self.current_frame + if end_frame is None: + end_frame = self.total_frames - 1 + + self.templates[template_id] = { + 'template': template.copy(), + 'region': region, + 'start_frame': start_frame, + 'end_frame': end_frame + } + + # Set as current template if it's the first one or if it covers current frame + if self.current_template_id is None or (start_frame <= self.current_frame <= end_frame): + self.current_template_id = template_id + self.tracking_template = template.copy() + self.template_region = region + + self.show_feedback_message(f"Template {template_id} added (frames {start_frame}-{end_frame})") + return template_id + + def remove_template(self, template_id): + """Remove a template from the collection""" + if template_id in self.templates: + del self.templates[template_id] + + # If we removed the current template, find a new one + if self.current_template_id == template_id: + self._select_best_template_for_frame(self.current_frame) + + self.show_feedback_message(f"Template {template_id} removed") + return True + return False + + def get_template_for_frame(self, frame_number): + """Get the best template for a given frame number""" + best_template_id = None + best_priority = -1 + + for template_id, template_data in self.templates.items(): + start_frame = template_data['start_frame'] + end_frame = template_data['end_frame'] + + # Check if template covers this frame + if start_frame <= frame_number <= end_frame: + # Priority: templates that start earlier and end later are preferred + priority = (end_frame - start_frame) - abs(frame_number - (start_frame + end_frame) // 2) + if priority > best_priority: + best_priority = priority + best_template_id = template_id + + return best_template_id + + def _select_best_template_for_frame(self, frame_number): + """Select the best template for the current frame""" + best_template_id = self.get_template_for_frame(frame_number) + + if best_template_id is not None: + self.current_template_id = best_template_id + template_data = self.templates[best_template_id] + self.tracking_template = template_data['template'].copy() + self.template_region = template_data['region'] + return True + else: + self.current_template_id = None + self.tracking_template = None + self.template_region = None + return False + + def navigate_to_next_template(self): + """Navigate to next template (; key)""" + if not self.templates: + return + + template_ids = sorted(self.templates.keys()) + if self.current_template_id is None: + # Find first template that starts after current frame + for template_id in template_ids: + if self.templates[template_id]['start_frame'] > self.current_frame: + self.current_template_id = template_id + self._select_best_template_for_frame(self.current_frame) + self.show_feedback_message(f"Template {template_id} selected") + return + else: + # Find next template + current_idx = template_ids.index(self.current_template_id) + for i in range(current_idx + 1, len(template_ids)): + template_id = template_ids[i] + if self.templates[template_id]['start_frame'] > self.current_frame: + self.current_template_id = template_id + self._select_best_template_for_frame(self.current_frame) + self.show_feedback_message(f"Template {template_id} selected") + return + + self.show_feedback_message("No next template found") + + def navigate_to_previous_template(self): + """Navigate to previous template (: key)""" + if not self.templates: + return + + template_ids = sorted(self.templates.keys()) + if self.current_template_id is None: + # Find last template that starts before current frame + for template_id in reversed(template_ids): + if self.templates[template_id]['start_frame'] < self.current_frame: + self.current_template_id = template_id + self._select_best_template_for_frame(self.current_frame) + self.show_feedback_message(f"Template {template_id} selected") + return + else: + # Find previous template + current_idx = template_ids.index(self.current_template_id) + for i in range(current_idx - 1, -1, -1): + template_id = template_ids[i] + if self.templates[template_id]['start_frame'] < self.current_frame: + self.current_template_id = template_id + self._select_best_template_for_frame(self.current_frame) + self.show_feedback_message(f"Template {template_id} selected") + return + + self.show_feedback_message("No previous template found") def apply_rotation(self, frame): """Apply rotation to frame""" @@ -2847,6 +3004,41 @@ class VideoEditor: 1, ) + # Draw template markers + for template_id, template_data in self.templates.items(): + start_frame = template_data['start_frame'] + end_frame = template_data['end_frame'] + + # Draw template range + start_progress = start_frame / max(1, self.total_frames - 1) + end_progress = end_frame / max(1, self.total_frames - 1) + start_x = bar_x_start + int(bar_width * start_progress) + end_x = bar_x_start + int(bar_width * end_progress) + + # Template range color (green for active, blue for inactive) + is_active = (self.current_template_id == template_id) + template_color = (0, 255, 0) if is_active else (255, 0, 0) # Green if active, red if inactive + + # Draw template range bar + cv2.rectangle( + frame, + (start_x, bar_y + 2), + (end_x, bar_y + self.TIMELINE_BAR_HEIGHT - 2), + template_color, + -1, + ) + + # Draw template ID number + cv2.putText( + frame, + str(template_id), + (start_x + 2, bar_y + 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.3, + (255, 255, 255), + 1, + ) + def display_current_frame(self): """Display the current frame with all overlays""" if self.current_display_frame is None: @@ -3309,6 +3501,21 @@ class VideoEditor: self.template_selection_start = None self.template_selection_rect = None + # Handle Ctrl+Right-click for template removal + if event == cv2.EVENT_RBUTTONDOWN and (flags & cv2.EVENT_FLAG_CTRLKEY) and not (flags & cv2.EVENT_FLAG_SHIFTKEY): + if not self.is_image_mode and self.templates: + # Check if click is near any template region + screen_x, screen_y = x, y + raw_x, raw_y = self._map_screen_to_rotated(screen_x, screen_y) + + for template_id, template_data in self.templates.items(): + tx, ty, tw, th = template_data['region'] + # Check if click is within template region + if (tx <= raw_x <= tx + tw and ty <= raw_y <= ty + th): + self.remove_template(template_id) + self.save_state() + return + # 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: @@ -4513,6 +4720,10 @@ class VideoEditor: print(f"DEBUG: Multi-scale template matching toggled to {self.multi_scale_template_matching}") self.show_feedback_message(f"Multi-scale template matching {'ON' if self.multi_scale_template_matching else 'OFF'}") self.save_state() + elif key == ord(";"): # Semicolon - Navigate to next template + self.navigate_to_next_template() + elif key == ord(":"): # Colon - Navigate to previous template + self.navigate_to_previous_template() elif key == ord("t"): # Marker looping only for videos if not self.is_image_mode: