shapes/_geometry.py

#!/usr/bin/env python3
"""
Generic polygon/ring pixel-sampling primitives shared by the per-shape
validators (diamond.py, circle.py, hexagon.py). No shape-specific thresholds
or dispatch logic lives here — just parameterized geometry sampling.
"""

import cv2
import numpy as np


def polygon_vertex_points(
    cx: int, cy: int, r: int, n_sides: int, angle_offset_deg: float = 0.0
) -> list[tuple[float, float]]:
    """
    Vertex positions of a regular n-gon centred at (cx, cy) with circumradius r.

    angle_offset_deg=-90 places the first vertex directly above the centre
    (image y-down convention), matching a diamond's top vertex.
    """
    pts = []
    for k in range(n_sides):
        theta = np.radians(angle_offset_deg + k * 360.0 / n_sides)
        pts.append((cx + r * np.cos(theta), cy + r * np.sin(theta)))
    return pts


def count_visible_sides(
    binary: np.ndarray,
    cx: int, cy: int, r: int,
    n_sides: int,
    angle_offset_deg: float = 0.0,
    n_samples: int = 8,
    hit_frac: float = 0.40,
) -> int:
    """
    Count how many of an n-gon's sides have visible (white) pixels along
    their interior, sampled at n_samples evenly-spaced points per side
    (avoiding the vertices themselves, t in [0.15, 0.85]).

    A side is 'present' if >= hit_frac of its samples hit a white pixel in a
    small 5x5 patch. Generalises the diamond-only visible-sides check to any
    regular polygon: diamond uses n_sides=4, angle_offset_deg=-90 (vertices at
    bbox midpoints — top/right/bottom/left); hexagon uses n_sides=6.
    """
    H, W = binary.shape
    pts = polygon_vertex_points(cx, cy, r, n_sides, angle_offset_deg)
    count = 0
    for i in range(n_sides):
        x1, y1 = pts[i]
        x2, y2 = pts[(i + 1) % n_sides]
        hits = 0
        for t in np.linspace(0.15, 0.85, n_samples):
            px = int(x1 + (x2 - x1) * t)
            py = int(y1 + (y2 - y1) * t)
            if 0 <= py < H and 0 <= px < W:
                patch = binary[max(0, py - 2):py + 3, max(0, px - 2):px + 3]
                if patch.any():
                    hits += 1
        if hits >= n_samples * hit_frac:
            count += 1
    return count


def count_hits_at_angles(
    binary: np.ndarray,
    cx: int, cy: int, r: int,
    angles_deg,
    dr_window: int = 6,
    n_probe: int = 5,
) -> int:
    """
    Count how many of the given angles (degrees, image y-down convention)
    have a white pixel near radius r.

    At each angle, probes n_probe points spaced across [r - dr_window,
    r + dr_window] and counts a hit if any of them lands on a white pixel
    in a small 3x3 patch. Generalises the diamond-only inter-vertex check
    (angles_deg=(45,135,225,315)) to arbitrary angle sets.
    """
    H, W = binary.shape
    count = 0
    for deg in angles_deg:
        rad = np.radians(deg)
        for dr in np.linspace(max(1, r - dr_window), r + dr_window, n_probe):
            py = int(cy + dr * np.sin(rad))
            px = int(cx + dr * np.cos(rad))
            if 0 <= py < H and 0 <= px < W:
                patch = binary[max(0, py - 3):py + 3, max(0, px - 3):px + 3]
                if patch.any():
                    count += 1
                    break
    return count


def ring_coverage_fraction(
    binary: np.ndarray,
    cx: int, cy: int, r: int,
    n_samples: int = 24,
) -> float:
    """
    Fraction of n_samples evenly-spaced angles around the full circle that
    have a white pixel near radius r. A real circle's boundary is continuous,
    so this should be high; a diamond's discrete 4 sides leave it low.
    """
    angles = np.linspace(0, 360, n_samples, endpoint=False)
    hits = count_hits_at_angles(binary, cx, cy, r, angles)
    return hits / n_samples if n_samples else 0.0


def outer_ring_density(
    binary: np.ndarray,
    cx: int, cy: int, r: int,
    width: int = 8,
) -> float:
    """
    Fraction of white pixels in the annular band just outside the estimated
    shape boundary. Real symbols float in open blueprint space (low density);
    table labels and grid-surrounded shapes have a high density.
    """
    H, W = binary.shape
    mask = np.zeros((H, W), np.uint8)
    cv2.circle(mask, (cx, cy), r + width, 255, thickness=-1)
    cv2.circle(mask, (cx, cy), max(1, r - 2), 0, thickness=-1)
    ring_pixels = int(cv2.countNonZero(mask))
    if ring_pixels == 0:
        return 0.0
    return float(cv2.countNonZero(cv2.bitwise_and(binary, binary, mask=mask))) / ring_pixels


def is_diamond_vertex_layout(
    approx: np.ndarray, x: int, y: int, w: int, h: int, vertex_tol: float = 0.22
) -> bool:
    """
    A diamond's 4 vertices sit near the MIDPOINTS of the bounding-box sides.
    An axis-aligned rectangle's vertices sit at the CORNERS.

    Each vertex must be within vertex_tol * min(w, h) of some expected midpoint.
    Used both by classify_shape() (to distinguish a diamond from a square
    among 4-vertex exemplars) and by diamond.passes_geometry().
    """
    pts = approx.reshape(-1, 2).astype(float)
    cx, cy = x + w / 2.0, y + h / 2.0

    # Canonical ordering by polar angle (independent of contour winding direction)
    order = np.argsort(np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx))
    pts = pts[order]

    midpoints = np.array([
        [x + w, cy   ],   # right  (≈  0°)
        [cx,    y + h],   # bottom (≈ 90°)
        [x,     cy   ],   # left   (≈180°)
        [cx,    y    ],   # top    (≈270°)
    ])

    tol = min(w, h) * vertex_tol
    for pt in pts:
        if np.min(np.linalg.norm(midpoints - pt, axis=1)) > tol:
            return False
    return True