Refactor template management in VideoEditor to simplify template structure
This commit updates the template management system in the VideoEditor by transitioning from a dictionary-based structure to a list of tuples. The new structure simplifies the handling of templates, focusing on start frames and regions without the need for template IDs or counters. Additionally, the loading and saving of templates have been streamlined, enhancing clarity and efficiency in template operations during video editing sessions.
This commit is contained in:
253
croppa/main.py
253
croppa/main.py
@@ -865,10 +865,8 @@ class VideoEditor:
|
|||||||
self.template_selection_start = None
|
self.template_selection_start = None
|
||||||
self.template_selection_rect = None
|
self.template_selection_rect = None
|
||||||
|
|
||||||
# Multiple templates system
|
# Simple template system - list of (start_frame, region) tuples sorted by start_frame
|
||||||
self.templates = {} # {template_id: {'template': image, 'region': (x,y,w,h), 'start_frame': int, 'end_frame': int}}
|
self.templates = [] # [(start_frame, region), ...] sorted by start_frame
|
||||||
self.current_template_id = None
|
|
||||||
self.template_id_counter = 0
|
|
||||||
|
|
||||||
# Template matching modes
|
# Template matching modes
|
||||||
self.template_matching_full_frame = False # Toggle for full frame vs cropped template matching
|
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()},
|
'tracking_points': {str(k): v for k, v in self.tracking_points.items()},
|
||||||
'feature_tracker': self.feature_tracker.get_state_dict(),
|
'feature_tracker': self.feature_tracker.get_state_dict(),
|
||||||
'template_matching_full_frame': self.template_matching_full_frame,
|
'template_matching_full_frame': self.template_matching_full_frame,
|
||||||
'templates': {str(k): {
|
'templates': [{
|
||||||
'template': None, # Don't save template images (too large)
|
'start_frame': start_frame,
|
||||||
'region': v['region'],
|
'region': region
|
||||||
'start_frame': v['start_frame'],
|
} for start_frame, region in self.templates]
|
||||||
'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:
|
with open(state_file, 'w') as f:
|
||||||
@@ -1019,22 +1013,14 @@ class VideoEditor:
|
|||||||
if 'template_matching_full_frame' in state:
|
if 'template_matching_full_frame' in state:
|
||||||
self.template_matching_full_frame = state['template_matching_full_frame']
|
self.template_matching_full_frame = state['template_matching_full_frame']
|
||||||
|
|
||||||
# Load multiple templates state
|
# Load simple templates state
|
||||||
if 'templates' in state:
|
if 'templates' in state:
|
||||||
self.templates = {}
|
self.templates = []
|
||||||
for template_id_str, template_data in state['templates'].items():
|
for template_data in state['templates']:
|
||||||
template_id = int(template_id_str)
|
self.templates.append((template_data['start_frame'], template_data['region']))
|
||||||
self.templates[template_id] = {
|
# Sort by start_frame
|
||||||
'template': None, # Will be recreated when needed
|
self.templates.sort(key=lambda x: x[0])
|
||||||
'region': template_data['region'],
|
|
||||||
'start_frame': template_data['start_frame'],
|
|
||||||
'end_frame': template_data['end_frame']
|
|
||||||
}
|
|
||||||
print(f"Loaded {len(self.templates)} templates")
|
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
|
# Validate cut markers against current video length
|
||||||
if self.cut_start_frame is not None and self.cut_start_frame >= self.total_frames:
|
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
|
# Calculate offset from template matching if enabled
|
||||||
template_offset = None
|
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.current_display_frame is not None:
|
||||||
if self.template_matching_full_frame:
|
if self.template_matching_full_frame:
|
||||||
# Full frame mode - use the entire original frame
|
# Full frame mode - use the entire original frame
|
||||||
@@ -1825,7 +1811,7 @@ class VideoEditor:
|
|||||||
|
|
||||||
def _get_template_matching_position(self, frame_number):
|
def _get_template_matching_position(self, frame_number):
|
||||||
"""Get template matching position and confidence for a frame"""
|
"""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
|
return None
|
||||||
|
|
||||||
if self.current_display_frame is not None:
|
if self.current_display_frame is not None:
|
||||||
@@ -2280,13 +2266,21 @@ class VideoEditor:
|
|||||||
|
|
||||||
def track_template(self, frame):
|
def track_template(self, frame):
|
||||||
"""Track the template in the current 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
|
return None
|
||||||
|
|
||||||
template_data = self.templates[self.current_template_id]
|
# Get the template for current frame
|
||||||
tracking_template = template_data['template']
|
template_index = self.get_template_for_frame(self.current_frame)
|
||||||
|
if template_index is None:
|
||||||
|
return None
|
||||||
|
|
||||||
if tracking_template is None:
|
start_frame, region = self.templates[template_index]
|
||||||
|
x, y, w, h = region
|
||||||
|
|
||||||
|
# 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
|
return None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -2393,137 +2387,86 @@ class VideoEditor:
|
|||||||
else:
|
else:
|
||||||
self.show_feedback_message("Template region outside frame bounds")
|
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"""
|
"""Add a new template to the collection"""
|
||||||
template_id = self.template_id_counter
|
|
||||||
self.template_id_counter += 1
|
|
||||||
|
|
||||||
if start_frame is None:
|
if start_frame is None:
|
||||||
start_frame = self.current_frame
|
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
|
# Add template to list
|
||||||
if self.current_template_id is not None and self.current_template_id in self.templates:
|
self.templates.append((start_frame, region))
|
||||||
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] = {
|
# Sort by start_frame
|
||||||
'template': template.copy(),
|
self.templates.sort(key=lambda x: x[0])
|
||||||
'region': region,
|
|
||||||
'start_frame': start_frame,
|
|
||||||
'end_frame': end_frame
|
|
||||||
}
|
|
||||||
|
|
||||||
# Set as current template
|
self.show_feedback_message(f"Template added at frame {start_frame}")
|
||||||
self.current_template_id = template_id
|
return len(self.templates) - 1
|
||||||
|
|
||||||
self.show_feedback_message(f"Template {template_id} added (frames {start_frame}-{end_frame})")
|
|
||||||
return template_id
|
|
||||||
|
|
||||||
def remove_template(self, template_id):
|
def remove_template(self, template_id):
|
||||||
"""Remove a template from the collection"""
|
"""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)
|
|
||||||
|
|
||||||
# If no templates left, clear current template
|
|
||||||
if not self.templates:
|
if not self.templates:
|
||||||
self.current_template_id = None
|
return False
|
||||||
|
|
||||||
self.show_feedback_message(f"Template {template_id} removed")
|
# Find template with start_frame > current_frame
|
||||||
|
current_frame = self.current_frame
|
||||||
|
template_to_remove = None
|
||||||
|
|
||||||
|
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 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
|
return True
|
||||||
|
else:
|
||||||
|
self.show_feedback_message("No template to remove")
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def get_template_for_frame(self, frame_number):
|
def get_template_for_frame(self, frame_number):
|
||||||
"""Get the most recent template that covers this frame"""
|
"""Get the template for the current frame"""
|
||||||
# Find the most recent template (highest ID) that covers this frame
|
if not self.templates:
|
||||||
best_template_id = None
|
return None
|
||||||
|
|
||||||
for template_id, template_data in self.templates.items():
|
# Find template with start_frame > current_frame
|
||||||
start_frame = template_data['start_frame']
|
for i, (start_frame, region) in enumerate(self.templates):
|
||||||
end_frame = template_data['end_frame']
|
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 None
|
||||||
|
|
||||||
# Check if template covers this frame
|
# If no template found with start_frame > current_frame, use the last one
|
||||||
if start_frame <= frame_number <= end_frame:
|
return len(self.templates) - 1 if self.templates else None
|
||||||
# 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):
|
def _select_best_template_for_frame(self, frame_number):
|
||||||
"""Select the most recent template for the current frame"""
|
"""Select the best template for the current frame"""
|
||||||
best_template_id = self.get_template_for_frame(frame_number)
|
template_index = self.get_template_for_frame(frame_number)
|
||||||
|
return template_index is not None
|
||||||
|
|
||||||
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
|
|
||||||
return True
|
|
||||||
else:
|
|
||||||
self.current_template_id = None
|
|
||||||
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']
|
|
||||||
|
|
||||||
# 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
|
|
||||||
else:
|
|
||||||
return False
|
|
||||||
else:
|
|
||||||
return False
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Error recreating template {template_id}: {e}")
|
|
||||||
return False
|
|
||||||
|
|
||||||
def jump_to_previous_template(self):
|
def jump_to_previous_template(self):
|
||||||
"""Jump to the previous template marker (frame where template was created)."""
|
"""Jump to the previous template marker (frame where template was created)."""
|
||||||
if self.is_image_mode:
|
if self.is_image_mode:
|
||||||
return
|
return
|
||||||
self.stop_auto_repeat_seek()
|
self.stop_auto_repeat_seek()
|
||||||
template_frames = sorted([data['start_frame'] for data in self.templates.values()])
|
if not self.templates:
|
||||||
if not template_frames:
|
|
||||||
print("DEBUG: No template markers; prev jump ignored")
|
print("DEBUG: No template markers; prev jump ignored")
|
||||||
return
|
return
|
||||||
current = self.current_frame
|
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:
|
if candidates:
|
||||||
target = candidates[-1]
|
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)
|
self.seek_to_frame(target)
|
||||||
else:
|
else:
|
||||||
target = template_frames[0]
|
target = self.templates[0][0]
|
||||||
print(f"DEBUG: Jump prev template to first marker from {current} -> {target}; template_frames={template_frames}")
|
print(f"DEBUG: Jump prev template to first marker from {current} -> {target}")
|
||||||
self.seek_to_frame(target)
|
self.seek_to_frame(target)
|
||||||
|
|
||||||
def jump_to_next_template(self):
|
def jump_to_next_template(self):
|
||||||
@@ -2531,18 +2474,17 @@ class VideoEditor:
|
|||||||
if self.is_image_mode:
|
if self.is_image_mode:
|
||||||
return
|
return
|
||||||
self.stop_auto_repeat_seek()
|
self.stop_auto_repeat_seek()
|
||||||
template_frames = sorted([data['start_frame'] for data in self.templates.values()])
|
if not self.templates:
|
||||||
if not template_frames:
|
|
||||||
print("DEBUG: No template markers; next jump ignored")
|
print("DEBUG: No template markers; next jump ignored")
|
||||||
return
|
return
|
||||||
current = self.current_frame
|
current = self.current_frame
|
||||||
for f in template_frames:
|
for start_frame, region in self.templates:
|
||||||
if f > current:
|
if start_frame > current:
|
||||||
print(f"DEBUG: Jump next template from {current} -> {f}; template_frames={template_frames}")
|
print(f"DEBUG: Jump next template from {current} -> {start_frame}")
|
||||||
self.seek_to_frame(f)
|
self.seek_to_frame(start_frame)
|
||||||
return
|
return
|
||||||
target = template_frames[-1]
|
target = self.templates[-1][0]
|
||||||
print(f"DEBUG: Jump next template to last marker from {current} -> {target}; template_frames={template_frames}")
|
print(f"DEBUG: Jump next template to last marker from {current} -> {target}")
|
||||||
self.seek_to_frame(target)
|
self.seek_to_frame(target)
|
||||||
|
|
||||||
def apply_rotation(self, frame):
|
def apply_rotation(self, frame):
|
||||||
@@ -3000,33 +2942,29 @@ class VideoEditor:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Draw template markers
|
# Draw template markers
|
||||||
for template_id, template_data in self.templates.items():
|
for start_frame, region in self.templates:
|
||||||
start_frame = template_data['start_frame']
|
# Draw template start point
|
||||||
end_frame = template_data['end_frame']
|
|
||||||
|
|
||||||
# Draw template range
|
|
||||||
start_progress = start_frame / max(1, self.total_frames - 1)
|
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)
|
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)
|
# Template color (green for active, red for inactive)
|
||||||
is_active = (self.current_template_id == template_id)
|
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
|
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(
|
cv2.rectangle(
|
||||||
frame,
|
frame,
|
||||||
(start_x, bar_y + 2),
|
(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,
|
template_color,
|
||||||
-1,
|
-1,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Draw template ID number
|
# Draw template number
|
||||||
cv2.putText(
|
cv2.putText(
|
||||||
frame,
|
frame,
|
||||||
str(template_id),
|
str(start_frame),
|
||||||
(start_x + 2, bar_y + 10),
|
(start_x + 2, bar_y + 10),
|
||||||
cv2.FONT_HERSHEY_SIMPLEX,
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
0.3,
|
0.3,
|
||||||
@@ -3245,8 +3183,7 @@ class VideoEditor:
|
|||||||
|
|
||||||
# Draw template matching point (blue circle with confidence)
|
# Draw template matching point (blue circle with confidence)
|
||||||
if (not self.is_image_mode and
|
if (not self.is_image_mode and
|
||||||
self.templates and
|
self.templates):
|
||||||
self.current_template_id is not None):
|
|
||||||
# Get template matching position for current frame
|
# Get template matching position for current frame
|
||||||
template_pos = self._get_template_matching_position(self.current_frame)
|
template_pos = self._get_template_matching_position(self.current_frame)
|
||||||
if template_pos:
|
if template_pos:
|
||||||
@@ -3484,17 +3421,15 @@ class VideoEditor:
|
|||||||
screen_x, screen_y = x, y
|
screen_x, screen_y = x, y
|
||||||
raw_x, raw_y = self._map_screen_to_rotated(screen_x, screen_y)
|
raw_x, raw_y = self._map_screen_to_rotated(screen_x, screen_y)
|
||||||
|
|
||||||
for template_id, template_data in self.templates.items():
|
for i, (start_frame, region) in enumerate(self.templates):
|
||||||
# Only check templates that cover current frame
|
tx, ty, tw, th = region
|
||||||
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_x = tx + tw // 2
|
||||||
center_y = ty + th // 2
|
center_y = ty + th // 2
|
||||||
|
|
||||||
# Check if click is within 10px of template center
|
# Check if click is within 10px of template center
|
||||||
distance = ((raw_x - center_x) ** 2 + (raw_y - center_y) ** 2) ** 0.5
|
distance = ((raw_x - center_x) ** 2 + (raw_y - center_y) ** 2) ** 0.5
|
||||||
if distance <= 10:
|
if distance <= 10:
|
||||||
self.remove_template(template_id)
|
self.remove_template(i) # Pass index instead of ID
|
||||||
self.save_state()
|
self.save_state()
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -3729,7 +3664,6 @@ class VideoEditor:
|
|||||||
|
|
||||||
# Reset templates
|
# Reset templates
|
||||||
self.templates.clear()
|
self.templates.clear()
|
||||||
self.current_template_id = None
|
|
||||||
|
|
||||||
# Reset cut markers
|
# Reset cut markers
|
||||||
self.cut_start_frame = None
|
self.cut_start_frame = None
|
||||||
@@ -4693,7 +4627,6 @@ class VideoEditor:
|
|||||||
# Clear all templates
|
# Clear all templates
|
||||||
if self.templates:
|
if self.templates:
|
||||||
self.templates.clear()
|
self.templates.clear()
|
||||||
self.current_template_id = None
|
|
||||||
print("DEBUG: All templates cleared")
|
print("DEBUG: All templates cleared")
|
||||||
self.show_feedback_message("All templates cleared")
|
self.show_feedback_message("All templates cleared")
|
||||||
else:
|
else:
|
||||||
|
Reference in New Issue
Block a user