main.py

import argparse
import logging
from pathlib import Path

import coloredlogs
import cv2 as cv
import supervision as sv

from modules import (
    EpisodeManager,
    ExerciseSetupError,
    ExerciseValidator,
    Table,
    TrackerState,
)
from modules.motion_manager import MotionManager
from pipeline import BilliardSystem, GlobalContext
from pipeline.stages import (
    AnnotationStage,
    BallDetectionStage,
    EpisodeTrackingStage,
    MovementStage,
    TableDetectionStage,
    TrackingStage,
    UndistortStage,
)
from strategies.annotation import (
    CircleAnnotatorStrategy,
    LabelAnnotatorStrategy,
    MotionAnnotatorStrategy,
)
from strategies.ball_detection import (
    YOLOBallDetectionStrategy,
)
from strategies.processing import (
    DisplayStrategy,
    SaveVideoStrategy,
    VideoProcessingRunner,
)
from strategies.tracking import SortTrackingStrategy

logger = logging.getLogger("billiard-system")
coloredlogs.DEFAULT_DATE_FORMAT = "%H:%M:%S"

VIDEO_PATH = Path("assets/!exercises/0002/success_2.mp4")
EMPTY_FRAME = cv.imread("assets/table_gray.png")

def parse_args():
    parser = argparse.ArgumentParser(description="Billiard System Video Processing")
    parser.add_argument(
        "--save", action="store_true", help="Save processed video to output directory"
    )
    parser.add_argument(
        "--start-frame",
        type=int,
        default=0,
        help="Frame number to start processing from (default: 60)",
    )
    parser.add_argument(
        "--output-name",
        type=str,
        help="Output filename for saved video (only valid when --save is used)",
    )
    parser.add_argument(
        "--log-file",
        action="store_true",
        help="Log to file instead of terminal (saves to billiard-system.log)",
    )
    parser.add_argument(
        "--exercise",
        type=str,
        help="Path to exercise JSON file for automatic validation",
    )

    args = parser.parse_args()

    if args.output_name and not args.save:
        parser.error("--output-name can only be used when --save is specified")

    return args


def main():
    args = parse_args()

    if args.log_file:
        logging.basicConfig(
            filename="billiard-system.log",
            filemode="w",
            encoding="utf-8",
            level=logging.DEBUG,
        )
    else:
        coloredlogs.install(
            level="INFO",
            fmt="%(asctime)s %(levelname)s %(name)s: %(message)s",
        )
    

    # downsampling = DownsamplingStrategy(enabled=True)
    # ball_detector = OpenCVBallDetectionStrategy(downsampling=downsampling)
    ball_detector = YOLOBallDetectionStrategy(min_confidence=0.5)

    table = Table(EMPTY_FRAME)
    tracker_state = TrackerState()
    motion_manager = MotionManager()

    context = GlobalContext(
        table=table,
        tracker_state=tracker_state,
        motion_manager=motion_manager,
        episode_manager=None,
        exercise_validator=None,
    )

    episode_manager = EpisodeManager(table=table, global_context=context)
    context.episode_manager = episode_manager

    exercise_validator = None
    if args.exercise:
        try:
            exercise_validator = ExerciseValidator(
                args.exercise, table, episode_manager
            )
            context.exercise_validator = exercise_validator
            logger.info(f"Exercise validation enabled: {args.exercise}")
        except ExerciseSetupError as e:
            logger.error(f"Failed to load exercise: {e}")
            return

    episode_manager.start()

    stages = [
        UndistortStage(),
        TableDetectionStage(),
        BallDetectionStage(strategy=ball_detector),
        TrackingStage(
            strategy=SortTrackingStrategy(
                max_age=6000,
                min_hits=10,
                iou_threshold=0.05,
                use_distance_matching=True,
                distance_weight=0.5,
                max_match_distance=250,
            )
        ),
        MovementStage(),
        EpisodeTrackingStage(),
        AnnotationStage(
            annotators=[
                CircleAnnotatorStrategy(),
                LabelAnnotatorStrategy(),
                MotionAnnotatorStrategy(),
            ]
        ),
    ]

    system = BilliardSystem(context, stages)

    video_info = sv.VideoInfo.from_video_path(video_path=VIDEO_PATH)
    frame_generator = sv.get_video_frames_generator(
        source_path=VIDEO_PATH, start=args.start_frame
    )

    if args.save:
        strategy = SaveVideoStrategy(
            video_path=VIDEO_PATH,
            video_info=video_info,
            start_frame=args.start_frame,
            output_name=args.output_name,
        )
    else:
        strategy = DisplayStrategy()

    processor = VideoProcessingRunner(strategy)
    processor.run(system, frame_generator)


if __name__ == "__main__":
    main()