preprocess_video.py

"""
Video Preprocessing Script for CursorTracker

This script converts video files into the required format for cursor tracking:
1. Extracts individual frames from video as PNG images
2. Generates background subtraction masks using MOG algorithm
3. Optionally discovers cursor templates using blob detection

Usage:
    python preprocess_video.py --video_path /path/to/video.mp4 --output_dir /path/to/output
    python preprocess_video.py --video_path /path/to/video.mp4 --output_dir /path/to/output --extract_templates
"""

import os
import sys
import argparse
import cv2
import numpy as np
import imageio
import skimage.morphology as morph
from tqdm import tqdm
from PIL import Image


def get_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Preprocess video for cursor tracking: extract frames and generate background masks.'
    )
    parser.add_argument(
        '--video_path',
        type=str,
        required=True,
        help='Path to input video file (e.g., video.mp4)'
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        required=True,
        help='Output directory where processed data will be saved'
    )
    parser.add_argument(
        '--extract_templates',
        action='store_true',
        help='Extract cursor templates using blob detection (optional)'
    )
    parser.add_argument(
        '--start_frame',
        type=int,
        default=0,
        help='Frame number to start processing from (default: 0)'
    )
    parser.add_argument(
        '--end_frame',
        type=int,
        default=None,
        help='Frame number to end processing at (default: process all frames)'
    )
    parser.add_argument(
        '--mog_history',
        type=int,
        default=1,
        help='Number of frames for MOG background subtractor history (default: 1)'
    )
    parser.add_argument(
        '--dilation_radius',
        type=int,
        default=5,
        help='Radius for morphological dilation of background masks (default: 5)'
    )
    return parser.parse_args()


def create_blob_detector(roi_size=(128, 128), blob_min_area=3,
                         blob_min_int=0.5, blob_max_int=0.95, blob_th_step=10):
    """
    Create a blob detector for cursor template discovery.

    Args:
        roi_size: Region of interest size (width, height)
        blob_min_area: Minimum area for blob detection
        blob_min_int: Minimum intensity threshold (0-1)
        blob_max_int: Maximum intensity threshold (0-1)
        blob_th_step: Threshold step for blob detection

    Returns:
        cv2.SimpleBlobDetector object
    """
    params = cv2.SimpleBlobDetector_Params()
    params.filterByArea = True
    params.minArea = blob_min_area
    params.maxArea = roi_size[0] * roi_size[1]
    params.filterByCircularity = False
    params.filterByColor = False
    params.filterByConvexity = False
    params.filterByInertia = False
    params.minThreshold = int(blob_min_int * 255)
    params.maxThreshold = int(blob_max_int * 255)
    params.thresholdStep = blob_th_step

    ver = cv2.__version__.split('.')
    if int(ver[0]) < 3:
        return cv2.SimpleBlobDetector(params)
    else:
        return cv2.SimpleBlobDetector_create(params)


def euclidean_distance(center1, center2):
    """Compute Euclidean distance between two points."""
    c1x, c1y = center1
    c2x, c2y = center2
    return np.sqrt((c1x - c2x)**2 + (c1y - c2y)**2)


def estimated_template_bbox(image, center_x, center_y, blob_diameter, allowance=3):
    """
    Compute bounding box for cursor template from blob information.

    Args:
        image: Input image
        center_x: Blob center x-coordinate
        center_y: Blob center y-coordinate
        blob_diameter: Diameter of detected blob
        allowance: Additional pixels to add around template

    Returns:
        Tuple of (top_left_x, top_left_y, bottom_right_x, bottom_right_y)
    """
    inscribed_square_length = np.sqrt(blob_diameter * blob_diameter / 2)
    top_left_x = max(0, center_x - inscribed_square_length / 2 - allowance)
    top_left_y = max(0, center_y - inscribed_square_length / 2 - allowance)
    bottom_right_x = min(center_x + inscribed_square_length / 2 + allowance, image.shape[1])
    bottom_right_y = min(center_y + inscribed_square_length / 2 + allowance, image.shape[0])

    return int(top_left_x), int(top_left_y), int(bottom_right_x), int(bottom_right_y)


def preprocess_video(video_path, output_dir, extract_templates=False,
                     start_frame=0, end_frame=None, mog_history=1, dilation_radius=5):
    """
    Main preprocessing function: extract frames and generate background masks.

    Args:
        video_path: Path to input video file
        output_dir: Directory to save processed data
        extract_templates: Whether to extract cursor templates
        start_frame: Starting frame number
        end_frame: Ending frame number (None = process all)
        mog_history: MOG background subtractor history parameter
        dilation_radius: Radius for morphological dilation
    """
    # Verify video file exists
    if not os.path.exists(video_path):
        raise FileNotFoundError(f"Video file not found: {video_path}")

    # Create output directories
    image_folder = os.path.join(output_dir, "images")
    background_folder = os.path.join(output_dir, "background")

    os.makedirs(image_folder, exist_ok=True)
    os.makedirs(background_folder, exist_ok=True)

    if extract_templates:
        template_folder = os.path.join(output_dir, "estimated_templates")
        os.makedirs(template_folder, exist_ok=True)
        blob_detector = create_blob_detector()
        template_sequence = []
        template_frames = []
        template_images = []
        consecutive_threshold = 5  # Number of consecutive frames for template extraction
        max_distance = 150  # Maximum pixel distance between consecutive detections

    # Open video file
    video = cv2.VideoCapture(video_path)
    if not video.isOpened():
        raise IOError(f"Could not open video file: {video_path}")

    # Get video properties
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS)
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

    if end_frame is None:
        end_frame = total_frames

    print(f"\nVideo Information:")
    print(f"  Path: {video_path}")
    print(f"  Resolution: {width}x{height}")
    print(f"  FPS: {fps}")
    print(f"  Total Frames: {total_frames}")
    print(f"  Processing frames: {start_frame} to {end_frame}")
    print(f"  Output directory: {output_dir}\n")

    # Initialize MOG background subtractor
    fgbg = cv2.bgsegm.createBackgroundSubtractorMOG(history=mog_history)

    frame_num = 0
    num_saved = 0

    # Process video frames
    pbar = tqdm(total=end_frame - start_frame, desc="Processing frames")

    while video.isOpened():
        ret, frame = video.read()

        if not ret or frame_num >= end_frame:
            break

        if frame_num >= start_frame:
            # Apply background subtraction
            bg_mask = fgbg.apply(frame)

            # Apply morphological dilation
            bg_mask_dilated = morph.dilation(bg_mask, footprint=morph.disk(dilation_radius))

            # Save frame
            frame_filename = f"{frame_num}.png"
            imageio.imsave(os.path.join(image_folder, frame_filename), frame)

            # Save background mask
            bg_filename = f"bgsub_{frame_num}.png"
            imageio.imsave(os.path.join(background_folder, bg_filename), bg_mask_dilated)

            # Extract templates if requested
            if extract_templates:
                keypoints = blob_detector.detect(bg_mask_dilated)

                if len(keypoints) == 1:
                    # Single blob detected - potential cursor
                    blob_center = keypoints[0].pt
                    blob_diameter = keypoints[0].size

                    # Check distance constraint with previous detection
                    if template_sequence:
                        prev_center = (template_sequence[-1][0], template_sequence[-1][1])
                        distance = euclidean_distance(prev_center, blob_center)

                        if distance <= max_distance:
                            template_sequence.append((blob_center[0], blob_center[1], blob_diameter))
                            template_frames.append(frame_num)
                            template_images.append(frame.copy())
                        else:
                            # Distance too large, reset sequence
                            template_sequence = [(blob_center[0], blob_center[1], blob_diameter)]
                            template_frames = [frame_num]
                            template_images = [frame.copy()]
                    else:
                        # Start new sequence
                        template_sequence = [(blob_center[0], blob_center[1], blob_diameter)]
                        template_frames = [frame_num]
                        template_images = [frame.copy()]

                    # Save template if we have enough consecutive detections
                    if len(template_sequence) == consecutive_threshold:
                        # Extract and save template from middle frame
                        mid_idx = len(template_sequence) // 2
                        center_x, center_y, diameter = template_sequence[mid_idx]
                        template_image = template_images[mid_idx]
                        template_frame = template_frames[mid_idx]

                        x1, y1, x2, y2 = estimated_template_bbox(
                            template_image, center_x, center_y, diameter
                        )

                        template_crop = template_image[y1:y2, x1:x2]
                        template_filename = f"frame_{template_frame}_template.png"
                        imageio.imsave(
                            os.path.join(template_folder, template_filename),
                            template_crop
                        )

                        # Reset sequence
                        template_sequence = []
                        template_frames = []
                        template_images = []
                else:
                    # Multiple or no blobs detected, reset sequence
                    template_sequence = []
                    template_frames = []
                    template_images = []

            num_saved += 1
            pbar.update(1)

        frame_num += 1

    pbar.close()
    video.release()

    print(f"\nProcessing complete!")
    print(f"  Frames saved: {num_saved}")
    print(f"  Images directory: {image_folder}")
    print(f"  Background masks directory: {background_folder}")

    if extract_templates:
        num_templates = len([f for f in os.listdir(template_folder) if f.endswith('.png')])
        print(f"  Templates extracted: {num_templates}")
        print(f"  Templates directory: {template_folder}")

    print(f"\nYou can now run cursor tracking with:")
    print(f"  python cursor_tracker_dp.py --video_name {os.path.basename(output_dir)} --base_dir {os.path.dirname(output_dir)}")


def main():
    """Main entry point."""
    args = get_args()

    try:
        preprocess_video(
            video_path=args.video_path,
            output_dir=args.output_dir,
            extract_templates=args.extract_templates,
            start_frame=args.start_frame,
            end_frame=args.end_frame,
            mog_history=args.mog_history,
            dilation_radius=args.dilation_radius
        )
    except Exception as e:
        print(f"\nError: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()