py-media-grader/croppa/tracking.py

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, x: int, y: int, radius: int = 50):
        """Remove a tracking point by frame and proximity to x,y"""
        if frame_number not in self.tracking_points:
            return False
            
        points = self.tracking_points[frame_number]
        for i, (px, py) in enumerate(points):
            # Calculate distance between points
            distance = ((px - x) ** 2 + (py - y) ** 2) ** 0.5
            if distance <= radius:
                del points[i]
                if not points:
                    del self.tracking_points[frame_number]
                return True
                
        return False
                
    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
        
        # If no base_zoom_center is provided, use the center of the crop rect
        if base_zoom_center is None and base_crop_rect is not None:
            x, y, w, h = base_crop_rect
            self.base_zoom_center = (x + w//2, y + h//2)
        else:
            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)