trigger-off/src/utils/action_utils.py

"""
Action / coordinate utility functions.

All coordinates are in normalised [0, 1] space.
"""

from __future__ import annotations

import re
from typing import Optional, Tuple


# ---------------------------------------------------------------------------
# Point parsing
# ---------------------------------------------------------------------------

_POINT_RE = re.compile(
    r"[\(\[]\s*"
    r"([+-]?\d+(?:\.\d+)?)"
    r"\s*[,\s]\s*"
    r"([+-]?\d+(?:\.\d+)?)"
    r"\s*[\)\]]"
)


def parse_point(pred_str: str) -> Optional[Tuple[float, float]]:
    """
    Parse a predicted (x, y) coordinate string.

    Handles formats:
      - "(0.5, 0.3)"
      - "[0.5 0.3]"
      - "(0.5,0.3)"
      - "(512, 384)" (absolute pixels — treated as-is, caller should normalise)

    Returns:
        (x, y) float tuple, or None if parsing fails.
    """
    if not pred_str:
        return None

    pred_str = pred_str.strip()

    # Try regex first
    m = _POINT_RE.search(pred_str)
    if m:
        x, y = float(m.group(1)), float(m.group(2))
        return x, y

    # Try comma/space separated without brackets
    parts = re.split(r"[,\s]+", pred_str.strip("()[] "))
    if len(parts) == 2:
        try:
            return float(parts[0]), float(parts[1])
        except ValueError:
            pass

    return None


# ---------------------------------------------------------------------------
# Bounding-box helpers
# ---------------------------------------------------------------------------

def point_in_bbox(
    point: Tuple[float, float],
    bbox: Tuple[float, float, float, float],
) -> bool:
    """
    Check whether a normalised point falls inside a normalised bounding box.

    Args:
        point: (x, y) in [0, 1].
        bbox: (x1, y1, x2, y2) in [0, 1] with x2 > x1, y2 > y1.

    Returns:
        True if point is inside or on the boundary of the bbox.
    """
    px, py = point
    x1, y1, x2, y2 = bbox
    return x1 <= px <= x2 and y1 <= py <= y2


def iou_click(
    pred_point: Tuple[float, float],
    gt_bbox: Tuple[float, float, float, float],
    click_radius: float = 0.05,
) -> bool:
    """
    Determine if a predicted click point "hits" the ground-truth bounding box.

    A click is considered a hit if:
      - It falls inside the gt_bbox, OR
      - It is within `click_radius` of the bbox centre (in Euclidean distance).

    Args:
        pred_point: Predicted (x, y) in [0, 1].
        gt_bbox: Ground-truth (x1, y1, x2, y2) in [0, 1].
        click_radius: Maximum allowed distance from bbox centre.

    Returns:
        True if the click is considered a hit.
    """
    # Inside the bbox
    if point_in_bbox(pred_point, gt_bbox):
        return True

    # Within radius of bbox centre
    x1, y1, x2, y2 = gt_bbox
    cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
    px, py = pred_point
    dist = ((px - cx) ** 2 + (py - cy) ** 2) ** 0.5
    return dist <= click_radius


# ---------------------------------------------------------------------------
# Action matching
# ---------------------------------------------------------------------------

def action_match(
    pred_str: str,
    gt_str: str,
    threshold: float = 0.05,
) -> bool:
    """
    Check whether a predicted action string matches the ground-truth.

    Both strings should be "(x, y)" coordinate strings.  Returns True if
    the Euclidean distance between the two parsed points is ≤ threshold.

    Args:
        pred_str: Predicted action string.
        gt_str: Ground-truth action string.
        threshold: Maximum Euclidean distance (in normalised coords).

    Returns:
        True if the actions match.
    """
    pred_pt = parse_point(pred_str)
    gt_pt = parse_point(gt_str)

    if pred_pt is None or gt_pt is None:
        return False

    dx = pred_pt[0] - gt_pt[0]
    dy = pred_pt[1] - gt_pt[1]
    dist = (dx ** 2 + dy ** 2) ** 0.5
    return dist <= threshold


def normalize_action_str(action_str: str) -> str:
    """
    Normalise an action string to a canonical "(x,y)" representation.

    Args:
        action_str: Raw action string such as "(0.5, 0.3)" or "[0.5 0.3]".

    Returns:
        Canonical string like "(0.5000,0.3000)", or the original string if
        parsing fails.
    """
    pt = parse_point(action_str)
    if pt is None:
        return action_str.strip()
    return f"({pt[0]:.4f},{pt[1]:.4f})"