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: