diff --git a/croppa/main.py b/croppa/main.py index e8aac79..436ab17 100644 --- a/croppa/main.py +++ b/croppa/main.py @@ -865,10 +865,8 @@ class VideoEditor: 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 + # Simple template system - list of (start_frame, region) tuples sorted by start_frame + self.templates = [] # [(start_frame, region), ...] sorted by start_frame # Template matching modes self.template_matching_full_frame = False # Toggle for full frame vs cropped template matching @@ -920,14 +918,10 @@ class VideoEditor: 'tracking_points': {str(k): v for k, v in self.tracking_points.items()}, 'feature_tracker': self.feature_tracker.get_state_dict(), 'template_matching_full_frame': self.template_matching_full_frame, - '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 + 'templates': [{ + 'start_frame': start_frame, + 'region': region + } for start_frame, region in self.templates] } with open(state_file, 'w') as f: @@ -1019,22 +1013,14 @@ class VideoEditor: if 'template_matching_full_frame' in state: self.template_matching_full_frame = state['template_matching_full_frame'] - # Load multiple templates state + # Load simple 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'] - } + self.templates = [] + for template_data in state['templates']: + self.templates.append((template_data['start_frame'], template_data['region'])) + # Sort by start_frame + self.templates.sort(key=lambda x: x[0]) 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: @@ -1655,7 +1641,7 @@ class VideoEditor: # Calculate offset from template matching if enabled template_offset = None - if self.templates and self.current_template_id is not None: + if self.templates: if self.current_display_frame is not None: if self.template_matching_full_frame: # Full frame mode - use the entire original frame @@ -1825,7 +1811,7 @@ class VideoEditor: def _get_template_matching_position(self, frame_number): """Get template matching position and confidence for a frame""" - if not self.templates or self.current_template_id is None: + if not self.templates: return None if self.current_display_frame is not None: @@ -2280,13 +2266,21 @@ class VideoEditor: def track_template(self, frame): """Track the template in the current frame""" - if self.current_template_id is None or self.current_template_id not in self.templates: + if not self.templates: return None + + # Get the template for current frame + template_index = self.get_template_for_frame(self.current_frame) + if template_index is None: + return None + + start_frame, region = self.templates[template_index] + x, y, w, h = region - template_data = self.templates[self.current_template_id] - tracking_template = template_data['template'] - - if tracking_template is None: + # Extract template from current frame + if (y + h <= frame.shape[0] and x + w <= frame.shape[1]): + tracking_template = frame[y:y+h, x:x+w] + else: return None try: @@ -2393,137 +2387,86 @@ class VideoEditor: else: self.show_feedback_message("Template region outside frame bounds") - def add_template(self, template, region, start_frame=None, end_frame=None): + def add_template(self, template, region, start_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 - # Only end the current template if it exists and starts at or before the current frame - if self.current_template_id is not None and self.current_template_id in self.templates: - current_template = self.templates[self.current_template_id] - if current_template['start_frame'] <= self.current_frame: - # End the current template at the previous frame - self.templates[self.current_template_id]['end_frame'] = self.current_frame - 1 - print(f"DEBUG: Ended current template {self.current_template_id} at frame {self.current_frame - 1}") - - self.templates[template_id] = { - 'template': template.copy(), - 'region': region, - 'start_frame': start_frame, - 'end_frame': end_frame - } + # Add template to list + self.templates.append((start_frame, region)) - # Set as current template - self.current_template_id = template_id + # Sort by start_frame + self.templates.sort(key=lambda x: x[0]) - self.show_feedback_message(f"Template {template_id} added (frames {start_frame}-{end_frame})") - return template_id + self.show_feedback_message(f"Template added at frame {start_frame}") + return len(self.templates) - 1 def remove_template(self, template_id): """Remove a template from the collection""" - if template_id in self.templates: - del self.templates[template_id] + if not self.templates: + return False - # 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) - - # If no templates left, clear current template - if not self.templates: - self.current_template_id = None - - self.show_feedback_message(f"Template {template_id} removed") - return True - return False - - def get_template_for_frame(self, frame_number): - """Get the most recent template that covers this frame""" - # Find the most recent template (highest ID) that covers this frame - best_template_id = None + # Find template with start_frame > current_frame + current_frame = self.current_frame + template_to_remove = None - 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: - # Use the template with the highest ID (most recent) - if best_template_id is None or template_id > best_template_id: - best_template_id = template_id - - return best_template_id - - def _select_best_template_for_frame(self, frame_number): - """Select the most recent template for the current frame""" - best_template_id = self.get_template_for_frame(frame_number) + for i, (start_frame, region) in enumerate(self.templates): + if start_frame > current_frame: + # Found template with start_frame > current_frame + # Remove the previous one (if it exists) + if i > 0: + template_to_remove = i - 1 + break - if best_template_id is not None: - self.current_template_id = best_template_id - template_data = self.templates[best_template_id] - - # Recreate template if it's None (loaded from state) - if template_data['template'] is None: - if self.current_display_frame is not None: - success = self._recreate_template_from_template_data(best_template_id, self.current_display_frame) - if not success: - return False - else: - return False - - # Template is already stored in the templates dict + if template_to_remove is not None: + removed_template = self.templates.pop(template_to_remove) + self.show_feedback_message(f"Template removed (was at frame {removed_template[0]})") return True else: - self.current_template_id = None + self.show_feedback_message("No template to remove") return False - def _recreate_template_from_template_data(self, template_id, frame): - """Recreate template from template data region""" - try: - template_data = self.templates[template_id] - x, y, w, h = template_data['region'] + def get_template_for_frame(self, frame_number): + """Get the template for the current frame""" + if not self.templates: + return None - # Ensure region is within frame bounds - if (x >= 0 and y >= 0 and - x + w <= frame.shape[1] and - y + h <= frame.shape[0]): - - # Extract template from frame - template = frame[y:y+h, x:x+w] - if template.size > 0: - self.templates[template_id]['template'] = template.copy() - return True + # Find template with start_frame > current_frame + for i, (start_frame, region) in enumerate(self.templates): + if start_frame > frame_number: + # Found template with start_frame > current_frame + # Use the previous one (if it exists) + if i > 0: + return i - 1 else: - return False - else: - return False - - except Exception as e: - print(f"Error recreating template {template_id}: {e}") - return False + return None + + # If no template found with start_frame > current_frame, use the last one + return len(self.templates) - 1 if self.templates else None + + def _select_best_template_for_frame(self, frame_number): + """Select the best template for the current frame""" + template_index = self.get_template_for_frame(frame_number) + return template_index is not None + def jump_to_previous_template(self): """Jump to the previous template marker (frame where template was created).""" if self.is_image_mode: return self.stop_auto_repeat_seek() - template_frames = sorted([data['start_frame'] for data in self.templates.values()]) - if not template_frames: + if not self.templates: print("DEBUG: No template markers; prev jump ignored") return current = self.current_frame - candidates = [f for f in template_frames if f < current] + candidates = [start_frame for start_frame, region in self.templates if start_frame < current] if candidates: target = candidates[-1] - print(f"DEBUG: Jump prev template from {current} -> {target}; template_frames={template_frames}") + print(f"DEBUG: Jump prev template from {current} -> {target}") self.seek_to_frame(target) else: - target = template_frames[0] - print(f"DEBUG: Jump prev template to first marker from {current} -> {target}; template_frames={template_frames}") + target = self.templates[0][0] + print(f"DEBUG: Jump prev template to first marker from {current} -> {target}") self.seek_to_frame(target) def jump_to_next_template(self): @@ -2531,18 +2474,17 @@ class VideoEditor: if self.is_image_mode: return self.stop_auto_repeat_seek() - template_frames = sorted([data['start_frame'] for data in self.templates.values()]) - if not template_frames: + if not self.templates: print("DEBUG: No template markers; next jump ignored") return current = self.current_frame - for f in template_frames: - if f > current: - print(f"DEBUG: Jump next template from {current} -> {f}; template_frames={template_frames}") - self.seek_to_frame(f) + for start_frame, region in self.templates: + if start_frame > current: + print(f"DEBUG: Jump next template from {current} -> {start_frame}") + self.seek_to_frame(start_frame) return - target = template_frames[-1] - print(f"DEBUG: Jump next template to last marker from {current} -> {target}; template_frames={template_frames}") + target = self.templates[-1][0] + print(f"DEBUG: Jump next template to last marker from {current} -> {target}") self.seek_to_frame(target) def apply_rotation(self, frame): @@ -3000,33 +2942,29 @@ class VideoEditor: ) # 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 + for start_frame, region in self.templates: + # Draw template start point 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 (green for active, red for inactive) + template_index = self.get_template_for_frame(self.current_frame) + is_active = (template_index is not None and self.templates[template_index] == (start_frame, region)) template_color = (0, 255, 0) if is_active else (255, 0, 0) # Green if active, red if inactive - # Draw template range bar + # Draw template start marker cv2.rectangle( frame, (start_x, bar_y + 2), - (end_x, bar_y + self.TIMELINE_BAR_HEIGHT - 2), + (start_x + 4, bar_y + self.TIMELINE_BAR_HEIGHT - 2), template_color, -1, ) - # Draw template ID number + # Draw template number cv2.putText( frame, - str(template_id), + str(start_frame), (start_x + 2, bar_y + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, @@ -3245,8 +3183,7 @@ class VideoEditor: # Draw template matching point (blue circle with confidence) if (not self.is_image_mode and - self.templates and - self.current_template_id is not None): + self.templates): # Get template matching position for current frame template_pos = self._get_template_matching_position(self.current_frame) if template_pos: @@ -3484,19 +3421,17 @@ class VideoEditor: 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(): - # Only check templates that cover current frame - if (template_data['start_frame'] <= self.current_frame <= template_data['end_frame']): - tx, ty, tw, th = template_data['region'] - center_x = tx + tw // 2 - center_y = ty + th // 2 - - # Check if click is within 10px of template center - distance = ((raw_x - center_x) ** 2 + (raw_y - center_y) ** 2) ** 0.5 - if distance <= 10: - self.remove_template(template_id) - self.save_state() - return + for i, (start_frame, region) in enumerate(self.templates): + tx, ty, tw, th = region + center_x = tx + tw // 2 + center_y = ty + th // 2 + + # Check if click is within 10px of template center + distance = ((raw_x - center_x) ** 2 + (raw_y - center_y) ** 2) ** 0.5 + if distance <= 10: + self.remove_template(i) # Pass index instead of ID + self.save_state() + return # Store tracking points in ROTATED frame coordinates (pre-crop) rx, ry = self._map_screen_to_rotated(x, y) @@ -3729,7 +3664,6 @@ class VideoEditor: # Reset templates self.templates.clear() - self.current_template_id = None # Reset cut markers self.cut_start_frame = None @@ -4693,7 +4627,6 @@ class VideoEditor: # Clear all templates if self.templates: self.templates.clear() - self.current_template_id = None print("DEBUG: All templates cleared") self.show_feedback_message("All templates cleared") else: