scripts/inference/table_structure.py

"""Classical CV table structure detection using line detection.

Detects horizontal and vertical lines in table images using morphological
operations, finds cell boundaries from line intersections, and returns
a structured grid.

This works better than TATR on engineering drawing BOM tables because:
- BOM tables have clear ruled lines (horizontal + vertical)
- Morphological line detection is robust to line thickness variation
- No ML model needed — pure geometry

Usage:
    from scripts.inference.table_structure import detect_table_structure
    cells = detect_table_structure(table_crop_bgr)
"""
from __future__ import annotations

import cv2
import numpy as np


def detect_table_structure(
    image: np.ndarray,
    min_line_length_ratio: float = 0.15,
    line_thickness: int = 2,
) -> dict | None:
    """Detect table rows and columns from ruled lines.

    Args:
        image: BGR table crop image.
        min_line_length_ratio: Minimum line length as fraction of image dimension.
        line_thickness: Morphological kernel size for line detection.

    Returns:
        dict with row_boundaries, col_boundaries, cell_grid or None if no structure found.
    """
    h, w = image.shape[:2]

    # Upscale small images
    if h < 400:
        scale = 400 / h
        image = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        h, w = image.shape[:2]

    # Convert to grayscale and binarize
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # Invert: lines become white on black
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # ── Detect horizontal lines ──
    h_kernel_len = max(int(w * min_line_length_ratio), 20)
    h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (h_kernel_len, 1))
    h_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_kernel, iterations=2)

    # ── Detect vertical lines ──
    v_kernel_len = max(int(h * min_line_length_ratio), 20)
    v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_kernel_len))
    v_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_kernel, iterations=2)

    # ── Extract line positions ──
    h_positions = _find_line_positions(h_lines, axis="horizontal")
    v_positions = _find_line_positions(v_lines, axis="vertical")

    if len(h_positions) < 2 or len(v_positions) < 2:
        return None

    # Add image boundaries if missing
    if h_positions[0] > h * 0.05:
        h_positions.insert(0, 0)
    if h_positions[-1] < h * 0.95:
        h_positions.append(h)
    if v_positions[0] > w * 0.05:
        v_positions.insert(0, 0)
    if v_positions[-1] < w * 0.95:
        v_positions.append(w)

    # Merge lines that are too close (within 5px)
    h_positions = _merge_close_positions(h_positions, threshold=max(5, h * 0.01))
    v_positions = _merge_close_positions(v_positions, threshold=max(5, w * 0.01))

    n_rows = len(h_positions) - 1
    n_cols = len(v_positions) - 1

    if n_rows < 1 or n_cols < 1:
        return None

    # Build cell grid
    cells = []
    for r in range(n_rows):
        for c in range(n_cols):
            y1 = int(h_positions[r])
            y2 = int(h_positions[r + 1])
            x1 = int(v_positions[c])
            x2 = int(v_positions[c + 1])

            # Skip very thin cells (likely line artifacts)
            if (y2 - y1) < 5 or (x2 - x1) < 5:
                continue

            cells.append({
                "row": r,
                "col": c,
                "bbox": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
            })

    return {
        "row_count": n_rows,
        "col_count": n_cols,
        "row_boundaries": [int(p) for p in h_positions],
        "col_boundaries": [int(p) for p in v_positions],
        "cells": cells,
        "image_size": {"width": w, "height": h},
    }


def _find_line_positions(line_mask: np.ndarray, axis: str) -> list[float]:
    """Find positions of detected lines by projecting onto an axis."""
    if axis == "horizontal":
        # Sum along x-axis to get row profile
        projection = np.sum(line_mask, axis=1)
    else:
        # Sum along y-axis to get column profile
        projection = np.sum(line_mask, axis=0)

    # Threshold: positions where projection is significant
    threshold = np.max(projection) * 0.3 if np.max(projection) > 0 else 0
    peaks = np.where(projection > threshold)[0]

    if len(peaks) == 0:
        return []

    # Cluster close peaks into single line positions
    positions = []
    cluster = [peaks[0]]
    for p in peaks[1:]:
        if p - cluster[-1] <= 3:
            cluster.append(p)
        else:
            positions.append(float(np.mean(cluster)))
            cluster = [p]
    positions.append(float(np.mean(cluster)))

    return positions


def _merge_close_positions(positions: list[float], threshold: float) -> list[float]:
    """Merge positions that are within threshold pixels of each other."""
    if not positions:
        return positions

    merged = [positions[0]]
    for p in positions[1:]:
        if p - merged[-1] > threshold:
            merged.append(p)
        else:
            # Replace with average
            merged[-1] = (merged[-1] + p) / 2
    return merged