visualize_results.py

"""
Visualization Script for CursorTracker

This script generates visualization images with bounding boxes drawn on frames
and creates an annotated video showing the tracked cursor.

Usage:
    python visualize_results.py --video_name video1 --base_dir /path/to/data
    python visualize_results.py --video_name video1 --base_dir /path/to/data --result_num 1
"""

import os
import sys
import argparse
import cv2
import numpy as np
import imageio
from tqdm import tqdm
from PIL import Image, ImageDraw, ImageFont


def get_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Visualize cursor tracking results with bounding boxes.'
    )
    parser.add_argument(
        '-v', '--video_name',
        type=str,
        required=True,
        help='Name of video directory (under base_dir)'
    )
    parser.add_argument(
        '--base_dir',
        type=str,
        default='./data/',
        help='Base directory containing video data (default: ./data/)'
    )
    parser.add_argument(
        '--result_num',
        type=int,
        default=None,
        help='Result folder number (e.g., 1 for our_results_1). If not specified, uses latest.'
    )
    parser.add_argument(
        '--bbox_color',
        type=str,
        default='0,255,0',
        help='Bounding box color in BGR format (default: 0,255,0 = green)'
    )
    parser.add_argument(
        '--bbox_thickness',
        type=int,
        default=2,
        help='Bounding box line thickness (default: 2)'
    )
    parser.add_argument(
        '--fps',
        type=int,
        default=30,
        help='Output video FPS (default: 30)'
    )
    parser.add_argument(
        '--quality',
        type=int,
        default=9,
        help='Output video quality 0-10, 10=best (default: 9)'
    )
    return parser.parse_args()


def load_tracking_results(results_file):
    """
    Load tracking results from file.

    Args:
        results_file: Path to tracking results file

    Returns:
        Dictionary mapping frame numbers to bounding boxes (x1, y1, x2, y2)
    """
    if not os.path.exists(results_file):
        raise FileNotFoundError(f"Results file not found: {results_file}")

    results = {}
    with open(results_file, 'r') as f:
        for line in f:
            parts = line.strip().split(',')
            if len(parts) >= 5:
                frame_num = parts[0]
                x1, y1, x2, y2 = int(parts[1]), int(parts[2]), int(parts[3]), int(parts[4])
                results[frame_num] = (x1, y1, x2, y2)

    return results


def draw_bounding_box(image, bbox, color=(0, 255, 0), thickness=2, show_coords=True):
    """
    Draw bounding box on image.

    Args:
        image: Input image (numpy array)
        bbox: Bounding box tuple (x1, y1, x2, y2)
        color: Box color in BGR format
        thickness: Line thickness
        show_coords: Whether to show coordinates as text

    Returns:
        Image with bounding box drawn
    """
    x1, y1, x2, y2 = bbox

    # Skip if bbox is at origin (no detection)
    if x1 == 0 and y1 == 0 and x2 == 0 and y2 == 0:
        return image

    # Draw rectangle
    cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)

    # Add text with coordinates if requested
    if show_coords:
        text = f"({x1},{y1})"
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.4
        font_thickness = 1

        # Get text size to draw background
        (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, font_thickness)

        # Draw text background
        cv2.rectangle(image,
                     (x1, y1 - text_height - baseline - 5),
                     (x1 + text_width, y1),
                     color, -1)

        # Draw text
        cv2.putText(image, text, (x1, y1 - baseline - 2),
                   font, font_scale, (0, 0, 0), font_thickness)

    return image


def generate_visualizations(video_name, base_dir, result_num=None,
                           bbox_color=(0, 255, 0), bbox_thickness=2,
                           fps=30, quality=9):
    """
    Generate visualization images and video.

    Args:
        video_name: Name of video directory
        base_dir: Base directory containing video data
        result_num: Result folder number (None = use latest)
        bbox_color: Bounding box color in BGR
        bbox_thickness: Box line thickness
        fps: Output video FPS
        quality: Video quality (0-10, 10=best)
    """
    # Set up directories
    video_dir = os.path.join(base_dir, video_name)
    image_folder = os.path.join(video_dir, "images")

    # Find results folder
    if result_num is None:
        # Find latest results folder
        results_folders = [f for f in os.listdir(video_dir) if f.startswith("our_results_")]
        if not results_folders:
            raise FileNotFoundError(f"No results folders found in {video_dir}")
        results_folders.sort(key=lambda x: int(x.split("_")[-1]))
        results_folder_name = results_folders[-1]
        result_num = int(results_folder_name.split("_")[-1])
    else:
        results_folder_name = f"our_results_{result_num}"

    results_folder = os.path.join(video_dir, results_folder_name)
    results_file = os.path.join(results_folder, "our_results.txt")

    # Create output directories
    viz_folder = os.path.join(results_folder, "visualizations")
    os.makedirs(viz_folder, exist_ok=True)

    print(f"\nVisualization Configuration:")
    print(f"  Video name: {video_name}")
    print(f"  Base directory: {base_dir}")
    print(f"  Results folder: {results_folder_name}")
    print(f"  Bounding box color: {bbox_color}")
    print(f"  Bounding box thickness: {bbox_thickness}")
    print(f"  Output FPS: {fps}")
    print(f"  Output quality: {quality}/10\n")

    # Verify directories exist
    if not os.path.exists(image_folder):
        raise FileNotFoundError(f"Images folder not found: {image_folder}")
    if not os.path.exists(results_folder):
        raise FileNotFoundError(f"Results folder not found: {results_folder}")

    # Load tracking results
    print(f"Loading tracking results from: {results_file}")
    tracking_results = load_tracking_results(results_file)
    print(f"  Loaded {len(tracking_results)} frame results")

    # Get sorted list of image files
    image_files = sorted([f for f in os.listdir(image_folder) if f.endswith('.png')],
                        key=lambda x: int(x.split('.')[0]))

    if len(image_files) == 0:
        raise ValueError(f"No PNG images found in {image_folder}")

    print(f"  Found {len(image_files)} frames to process\n")

    # Process frames
    print("Generating visualization images...")
    video_frames = []

    for image_file in tqdm(image_files, desc="Processing frames"):
        frame_num = image_file.split('.')[0]
        image_path = os.path.join(image_folder, image_file)

        # Read image
        image = cv2.imread(image_path)
        if image is None:
            print(f"Warning: Could not read {image_path}")
            continue

        # Draw bounding box if available
        if frame_num in tracking_results:
            bbox = tracking_results[frame_num]
            image = draw_bounding_box(image, bbox, bbox_color, bbox_thickness)

        # Save visualization image
        viz_filename = f"viz_{image_file}"
        viz_path = os.path.join(viz_folder, viz_filename)
        cv2.imwrite(viz_path, image)

        # Add to video frames
        video_frames.append(image)

    # Create video
    if len(video_frames) > 0:
        print(f"\nCreating annotated video...")
        video_path = os.path.join(results_folder, f"tracked_video_{results_folder_name}.mp4")

        # Get frame dimensions
        height, width = video_frames[0].shape[:2]

        # Create video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_path, fourcc, fps, (width, height))

        for frame in tqdm(video_frames, desc="Writing video"):
            out.write(frame)

        out.release()

        print(f"\n✓ Visualization complete!")
        print(f"  Visualization images: {viz_folder}/")
        print(f"  Annotated video: {video_path}")
        print(f"  Total frames: {len(video_frames)}")

        # Calculate statistics
        detected_frames = sum(1 for frame_num in tracking_results.values()
                             if not (frame_num == (0, 0, 0, 0)))
        total_frames = len(image_files)
        detection_rate = (detected_frames / total_frames * 100) if total_frames > 0 else 0

        print(f"\nTracking Statistics:")
        print(f"  Frames with detections: {detected_frames}/{total_frames} ({detection_rate:.1f}%)")
    else:
        print("\nError: No frames were processed successfully")


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

    # Parse bbox color
    bbox_color = tuple(map(int, args.bbox_color.split(',')))

    try:
        generate_visualizations(
            video_name=args.video_name,
            base_dir=args.base_dir,
            result_num=args.result_num,
            bbox_color=bbox_color,
            bbox_thickness=args.bbox_thickness,
            fps=args.fps,
            quality=args.quality
        )
    except Exception as e:
        print(f"\nError: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()