drs_modules/trajectory.py

"""Trajectory estimation for the cricket ball.

Professional ball tracking systems reconstruct the ball's path in 3D
from several camera angles and then use physics or machine learning
models to project its flight.  Here we implement a far simpler
approach.  Given a sequence of ball centre coordinates extracted from
a single camera (behind the bowler), we fit a polynomial curve to
approximate the ball's trajectory in image space.  We assume that the
ball travels roughly along a parabolic path, so a quadratic fit to
``y`` as a function of ``x`` is appropriate for the vertical drop.

Because we lack explicit knowledge of the camera's field of view, the
stumps' location is estimated relative to the range of observed ball
positions.  If the projected path intersects a fixed region near the
bottom middle of the frame, we say that the ball would have hit the
stumps.
"""

from __future__ import annotations

import numpy as np
from typing import Callable, Dict, List, Tuple


def estimate_trajectory(centers: List[Tuple[int, int]], timestamps: List[float]) -> Dict[str, object]:
    """Fit a polynomial to the ball's path.

    Parameters
    ----------
    centers: list of tuple(int, int)
        Detected ball centre positions in pixel coordinates (x, y).
    timestamps: list of float
        Timestamps (in seconds) corresponding to each detection.  Unused
        in the current implementation but retained for extensibility.

    Returns
    -------
    dict
        A dictionary with keys ``coeffs`` (the polynomial coefficients
        [a, b, c] for ``y = a*x^2 + b*x + c``) and ``model`` (a
        callable that accepts an x coordinate and returns the
        predicted y coordinate).
    """
    if not centers:
        # No detections; return a dummy model
        return {"coeffs": np.array([0.0, 0.0, 0.0]), "model": lambda x: 0 * x}

    xs = np.array([pt[0] for pt in centers], dtype=np.float64)
    ys = np.array([pt[1] for pt in centers], dtype=np.float64)

    # Require at least 3 points for a quadratic fit; otherwise fall back
    # to a linear fit
    if len(xs) >= 3:
        coeffs = np.polyfit(xs, ys, 2)
        def model(x: np.ndarray | float) -> np.ndarray | float:
            return coeffs[0] * (x ** 2) + coeffs[1] * x + coeffs[2]
    else:
        coeffs = np.polyfit(xs, ys, 1)
        def model(x: np.ndarray | float) -> np.ndarray | float:
            return coeffs[0] * x + coeffs[1]

    return {"coeffs": coeffs, "model": model}


def predict_stumps_intersection(trajectory: Dict[str, object]) -> bool:
    """Predict whether the ball's trajectory will hit the stumps.

    The stumps are assumed to lie roughly in the centre of the frame
    along the horizontal axis and occupy the lower quarter of the
    vertical axis.  This heuristic works reasonably well for videos
    captured from behind the bowler.  In a production system you
    would calibrate the exact position of the stumps from the pitch
    geometry.

    Parameters
    ----------
    trajectory: dict
        Output of :func:`estimate_trajectory`, containing the
        polynomial model and the original ``centers`` list if needed.

    Returns
    -------
    bool
        True if the ball is predicted to hit the stumps, False otherwise.
    """
    model: Callable[[float], float] = trajectory["model"]
    coeffs = trajectory["coeffs"]

    # Recover approximate frame dimensions from the observed centres.  We
    # estimate the width and height as slightly larger than the max
    # observed coordinates.
    # Note: trajectory does not contain the centres directly, so we
    # recompute width and height heuristically based on coefficient
    # magnitudes.  To avoid overcomplication we assign reasonable
    # defaults if no centres were available.
    if hasattr(trajectory, "centers"):
        # never executed; left as placeholder
        pass

    # Use coefficients to infer approximate domain of x.  The roots of
    # derivative give extremum; but we simply sample across a range
    # derived from typical video width (e.g. 640px)
    frame_width = 640
    frame_height = 360

    # Estimate ball y position at the x coordinate corresponding to the
    # middle stump: 50% of frame width
    stumps_x = frame_width * 0.5
    predicted_y = model(stumps_x)

    # Define the vertical bounds of the wicket region in pixels.  The
    # top of the stumps is roughly three quarters down the frame and
    # the bottom is at the very bottom.  These ratios can be tuned.
    stumps_y_low = frame_height * 0.65
    stumps_y_high = frame_height * 0.95

    return stumps_y_low <= predicted_y <= stumps_y_high