150 lines
6.3 KiB
Python
150 lines
6.3 KiB
Python
from typing import List, Dict, Tuple, Optional
|
|
|
|
|
|
class MotionTracker:
|
|
"""Handles motion tracking for crop and pan operations"""
|
|
|
|
def __init__(self):
|
|
self.tracking_points = {} # {frame_number: [(x, y), ...]}
|
|
self.tracking_enabled = False
|
|
self.base_crop_rect = None # Original crop rect when tracking started
|
|
self.base_zoom_center = None # Original zoom center when tracking started
|
|
|
|
def add_tracking_point(self, frame_number: int, x: int, y: int):
|
|
"""Add a tracking point at the specified frame and coordinates"""
|
|
if frame_number not in self.tracking_points:
|
|
self.tracking_points[frame_number] = []
|
|
self.tracking_points[frame_number].append((x, y))
|
|
|
|
def remove_tracking_point(self, frame_number: int, point_index: int):
|
|
"""Remove a tracking point by frame and index"""
|
|
if frame_number in self.tracking_points and 0 <= point_index < len(self.tracking_points[frame_number]):
|
|
del self.tracking_points[frame_number][point_index]
|
|
if not self.tracking_points[frame_number]:
|
|
del self.tracking_points[frame_number]
|
|
|
|
def clear_tracking_points(self):
|
|
"""Clear all tracking points"""
|
|
self.tracking_points.clear()
|
|
|
|
def get_tracking_points_for_frame(self, frame_number: int) -> List[Tuple[int, int]]:
|
|
"""Get all tracking points for a specific frame"""
|
|
return self.tracking_points.get(frame_number, [])
|
|
|
|
def has_tracking_points(self) -> bool:
|
|
"""Check if any tracking points exist"""
|
|
return bool(self.tracking_points)
|
|
|
|
def get_interpolated_position(self, frame_number: int) -> Optional[Tuple[float, float]]:
|
|
"""Get interpolated position for a frame based on tracking points"""
|
|
if not self.tracking_points:
|
|
return None
|
|
|
|
# Get all frames with tracking points
|
|
frames = sorted(self.tracking_points.keys())
|
|
|
|
if not frames:
|
|
return None
|
|
|
|
# If we have a point at this exact frame, return it
|
|
if frame_number in self.tracking_points:
|
|
points = self.tracking_points[frame_number]
|
|
if points:
|
|
# Return average of all points at this frame
|
|
avg_x = sum(p[0] for p in points) / len(points)
|
|
avg_y = sum(p[1] for p in points) / len(points)
|
|
return (avg_x, avg_y)
|
|
|
|
# If frame is before first tracking point
|
|
if frame_number < frames[0]:
|
|
points = self.tracking_points[frames[0]]
|
|
if points:
|
|
avg_x = sum(p[0] for p in points) / len(points)
|
|
avg_y = sum(p[1] for p in points) / len(points)
|
|
return (avg_x, avg_y)
|
|
|
|
# If frame is after last tracking point
|
|
if frame_number > frames[-1]:
|
|
points = self.tracking_points[frames[-1]]
|
|
if points:
|
|
avg_x = sum(p[0] for p in points) / len(points)
|
|
avg_y = sum(p[1] for p in points) / len(points)
|
|
return (avg_x, avg_y)
|
|
|
|
# Find the two frames to interpolate between
|
|
for i in range(len(frames) - 1):
|
|
if frames[i] <= frame_number <= frames[i + 1]:
|
|
frame1, frame2 = frames[i], frames[i + 1]
|
|
points1 = self.tracking_points[frame1]
|
|
points2 = self.tracking_points[frame2]
|
|
|
|
if not points1 or not points2:
|
|
continue
|
|
|
|
# Get average positions for each frame
|
|
avg_x1 = sum(p[0] for p in points1) / len(points1)
|
|
avg_y1 = sum(p[1] for p in points1) / len(points1)
|
|
avg_x2 = sum(p[0] for p in points2) / len(points2)
|
|
avg_y2 = sum(p[1] for p in points2) / len(points2)
|
|
|
|
# Linear interpolation
|
|
t = (frame_number - frame1) / (frame2 - frame1)
|
|
interp_x = avg_x1 + t * (avg_x2 - avg_x1)
|
|
interp_y = avg_y1 + t * (avg_y2 - avg_y1)
|
|
|
|
return (interp_x, interp_y)
|
|
|
|
return None
|
|
|
|
def get_tracking_offset(self, frame_number: int) -> Tuple[float, float]:
|
|
"""Get the offset to center the crop on the tracked point"""
|
|
if not self.tracking_enabled or not self.base_zoom_center:
|
|
return (0.0, 0.0)
|
|
|
|
current_pos = self.get_interpolated_position(frame_number)
|
|
if not current_pos:
|
|
return (0.0, 0.0)
|
|
|
|
# Calculate offset to center the crop on the tracked point
|
|
# The offset should move the display so the tracked point stays centered
|
|
offset_x = current_pos[0] - self.base_zoom_center[0]
|
|
offset_y = current_pos[1] - self.base_zoom_center[1]
|
|
|
|
return (offset_x, offset_y)
|
|
|
|
def start_tracking(self, base_crop_rect: Tuple[int, int, int, int], base_zoom_center: Tuple[int, int]):
|
|
"""Start motion tracking with base positions"""
|
|
self.tracking_enabled = True
|
|
self.base_crop_rect = base_crop_rect
|
|
self.base_zoom_center = base_zoom_center
|
|
|
|
def stop_tracking(self):
|
|
"""Stop motion tracking"""
|
|
self.tracking_enabled = False
|
|
self.base_crop_rect = None
|
|
self.base_zoom_center = None
|
|
|
|
def to_dict(self) -> Dict:
|
|
"""Convert to dictionary for serialization"""
|
|
return {
|
|
'tracking_points': self.tracking_points,
|
|
'tracking_enabled': self.tracking_enabled,
|
|
'base_crop_rect': self.base_crop_rect,
|
|
'base_zoom_center': self.base_zoom_center
|
|
}
|
|
|
|
def from_dict(self, data: Dict):
|
|
"""Load from dictionary for deserialization"""
|
|
# Convert string keys back to integers for tracking_points
|
|
tracking_points_data = data.get('tracking_points', {})
|
|
self.tracking_points = {}
|
|
for frame_str, points in tracking_points_data.items():
|
|
frame_num = int(frame_str) # Convert string key to integer
|
|
self.tracking_points[frame_num] = points
|
|
|
|
self.tracking_enabled = data.get('tracking_enabled', False)
|
|
self.base_crop_rect = data.get('base_crop_rect', None)
|
|
self.base_zoom_center = data.get('base_zoom_center', None)
|
|
|
|
|