File size: 5,230 Bytes
2ebf7ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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
|