120class GridDetector:
121    """
122    Detects table grid intersections using morphological filtering and template matching.
123
124    This detector implements a multi-stage pipeline:
125
126    1. **Binarization**: Sauvola adaptive thresholding to handle varying lighting
127    2. **Morphological operations**: Dilation to connect broken rule segments
128    3. **Cross-kernel matching**: Template matching with a cross-shaped kernel to find
129       rule intersections where horizontal and vertical lines meet
130    4. **Grid growing**: Iterative point detection starting from a known seed point
131
132    The cross-kernel is designed to match the specific geometry of your table rules.
133    It should be sized so that after morphology, it aligns with actual corner shapes.
134
135    ## Tuning Guidelines
136
137    - **kernel_size**: Increase if you need more selectivity (fewer false positives)
138    - **cross_width/height**: Should match rule thickness after morphology
139    - **morph_size**: Increase to connect more broken lines, but this thickens rules
140    - **sauvola_k**: Increase to threshold more aggressively (remove noise)
141    - **search_region**: Increase for documents with more warping/distortion
142    - **distance_penalty**: Increase to prefer corners closer to expected positions
143
144    ## Visual Debugging
145
146    Set `visual=True` in methods to see intermediate results and tune parameters.
147    """
148
149    def __init__(
150        self,
151        kernel_size: int = 21,
152        cross_width: int = 6,
153        cross_height: Optional[int] = None,
154        morph_size: Optional[int] = None,
155        sauvola_k: float = 0.04,
156        sauvola_window: int = 15,
157        scale: float = 1.0,
158        search_region: int = 40,
159        distance_penalty: float = 0.4,
160        skip_astar_threshold: float = 0.2,
161        min_rows: int = 5,
162        grow_threshold: float = 0.3,
163        look_distance: int = 4,
164        cuts: int = 3,
165        cut_fraction: float = 0.5,
166    ):
167        """
168        Args:
169            kernel_size (int): the size of the cross kernel
170                a larger kernel size often means that more penalty is applied, often leading
171                to more sparse results
172            cross_width (int): the width of one of the edges in the cross filter, should be
173                roughly equal to the width of the rules in the image after morphology is applied
174            cross_height (int | None): useful if the horizontal rules and vertical rules
175                have different sizes
176            morph_size (int | None): the size of the morphology operators that are applied before
177                the cross kernel. 'bridges the gaps' of broken-up lines
178            sauvola_k (float): threshold parameter for sauvola thresholding
179            sauvola_window (int): window_size parameter for sauvola thresholding
180            scale (float): image scale factor to do calculations on (useful for increasing calculation speed mostly)
181            search_region (int): area in which to search for a new max value in `find_nearest` etc.
182            distance_penalty (float): how much the point finding algorithm penalizes points that are further in the region [0, 1]
183            skip_astar_threshold (float): minimum confidence score during table growing based on heuristic jump on which to skipastar pathfinding
184            min_rows (int): minimum number of rows to find before stopping the table finding algorithm
185            grow_threshold (float): the threshold for accepting a new point when growing the table
186            look_distance (int): how many points away to look when calculating the median slope
187            cuts (int): The amount of cuts (large deletions) to do in the grid during table growing
188            cut_fraction (float): The portion of the already-chosen corner points to delete during cutting
189        """
190        self._validate_parameters(
191            kernel_size,
192            cross_width,
193            cross_height,
194            morph_size,
195            search_region,
196            sauvola_k,
197            sauvola_window,
198            distance_penalty,
199            skip_astar_threshold,
200            cuts,
201            cut_fraction,
202        )
203
204        self._kernel_size = kernel_size
205        self._cross_width = cross_width
206        self._cross_height = cross_width if cross_height is None else cross_height
207        self._morph_size = morph_size if morph_size is not None else cross_width
208        self._search_region = search_region
209        self._sauvola_k = sauvola_k
210        self._sauvola_window = sauvola_window
211        self._distance_penalty = distance_penalty
212        self._scale = scale
213        self._skip_astar_threshold = skip_astar_threshold
214        self._min_rows = min_rows
215        self._grow_threshold = grow_threshold
216        self._look_distance = look_distance
217        self._cuts = cuts
218        self._cut_fraction = cut_fraction
219
220        self._cross_kernel = self._create_cross_kernel()
221
222    def _validate_parameters(
223        self,
224        kernel_size: int,
225        cross_width: int,
226        cross_height: Optional[int],
227        morph_size: Optional[int],
228        search_region: int,
229        sauvola_k: float,
230        sauvola_window: int,
231        distance_penalty: float,
232        skip_astar_threshold: float,
233        cuts: int,
234        cut_fraction: float,
235    ) -> None:
236        """Validate initialization parameters."""
237        if kernel_size % 2 == 0:
238            raise ValueError("kernel_size must be odd")
239        if (
240            kernel_size <= 0
241            or cross_width <= 0
242            or search_region <= 0
243            or sauvola_window <= 0
244        ):
245            raise ValueError("Size parameters must be positive")
246        if cross_height is not None and cross_height <= 0:
247            raise ValueError("cross_height must be positive")
248        if morph_size is not None and morph_size <= 0:
249            raise ValueError("morph_size must be positive")
250        if not 0 <= distance_penalty <= 1:
251            raise ValueError("distance_penalty must be in [0, 1]")
252        if sauvola_k <= 0:
253            raise ValueError("sauvola_k must be positive")
254        if skip_astar_threshold < 0 or skip_astar_threshold > 1:
255            raise ValueError("skip_astar_threshold must be in [0, 1]")
256        if cut_fraction < 0 or cut_fraction > 1:
257            raise ValueError("cut_fraction must be in [0, 1]")
258        if cuts < 0:
259            raise ValueError("cuts must be zero or positive")
260
261    def _create_gaussian_weights(self, region_size: int) -> NDArray:
262        """
263        Create a 2D Gaussian weight mask.
264
265        Args:
266            shape (tuple[int, int]): Shape of the region (height, width)
267            p (float): Minimum value at the edge = 1 - p
268
269        Returns:
270            NDArray: Gaussian weight mask
271        """
272        if self._distance_penalty == 0:
273            return np.ones((region_size, region_size), dtype=np.float32)
274
275        y = np.linspace(-1, 1, region_size)
276        x = np.linspace(-1, 1, region_size)
277        xv, yv = np.meshgrid(x, y)
278        dist_squared = xv**2 + yv**2
279
280        # Prevent log(0) when distance_penalty is 1
281        if self._distance_penalty >= 0.999:
282            sigma = 0.1  # Small sigma for very sharp peak
283        else:
284            sigma = np.sqrt(-1 / (2 * np.log(1 - self._distance_penalty)))
285
286        weights = np.exp(-dist_squared / (2 * sigma**2))
287
288        return weights.astype(np.float32)
289
290    def _create_cross_kernel(self) -> NDArray:
291        kernel = np.zeros((self._kernel_size, self._kernel_size), dtype=np.uint8)
292        center = self._kernel_size // 2
293
294        # Create horizontal bar
295        h_start = max(0, center - self._cross_height // 2)
296        h_end = min(self._kernel_size, center + (self._cross_height + 1) // 2)
297        kernel[h_start:h_end, :] = 255
298
299        # Create vertical bar
300        v_start = max(0, center - self._cross_width // 2)
301        v_end = min(self._kernel_size, center + (self._cross_width + 1) // 2)
302        kernel[:, v_start:v_end] = 255
303
304        return kernel
305
306    def _apply_morphology(self, binary: MatLike) -> MatLike:
307        # Define a horizontal kernel (adjust width as needed)
308        kernel_hor = cv.getStructuringElement(cv.MORPH_RECT, (self._morph_size, 1))
309        kernel_ver = cv.getStructuringElement(cv.MORPH_RECT, (1, self._morph_size))
310
311        # Apply dilation
312        dilated = cv.dilate(binary, kernel_hor, iterations=1)
313        dilated = cv.dilate(dilated, kernel_ver, iterations=1)
314
315        return dilated
316
317    def _apply_cross_matching(self, img: MatLike) -> MatLike:
318        """Apply cross kernel template matching."""
319        pad_y = self._cross_kernel.shape[0] // 2
320        pad_x = self._cross_kernel.shape[1] // 2
321
322        padded = cv.copyMakeBorder(
323            img, pad_y, pad_y, pad_x, pad_x, borderType=cv.BORDER_CONSTANT, value=0
324        )
325
326        filtered = cv.matchTemplate(padded, self._cross_kernel, cv.TM_SQDIFF_NORMED)
327        # Invert and normalize to 0-255 range
328        filtered = cv.normalize(1.0 - filtered, None, 0, 255, cv.NORM_MINMAX)
329        return filtered.astype(np.uint8)
330
331    def apply(self, img: MatLike, visual: bool = False) -> MatLike:
332        """
333        Apply the grid detection filter to the input image.
334
335        Args:
336            img (MatLike): the input image
337            visual (bool): whether to show intermediate steps
338
339        Returns:
340            MatLike: the filtered image, with high values (whiter pixels) at intersections of horizontal and vertical rules
341        """
342
343        if img is None or img.size == 0:
344            raise ValueError("Input image is empty or None")
345
346        binary = imu.sauvola(img, k=self._sauvola_k, window_size=self._sauvola_window)
347
348        if visual:
349            imu.show(binary, title="thresholded")
350
351        binary = self._apply_morphology(binary)
352
353        if visual:
354            imu.show(binary, title="dilated")
355
356        filtered = self._apply_cross_matching(binary)
357
358        return filtered
359
360    @log_calls(level=logging.DEBUG, include_return=True)
361    def find_nearest(
362        self, filtered: MatLike, point: Point, region: Optional[int] = None
363    ) -> Tuple[Point, float]:
364        """
365        Find the nearest 'corner match' in the image, along with its score [0,1]
366
367        Args:
368            filtered (MatLike): the filtered image (obtained through `apply`)
369            point (tuple[int, int]): the approximate target point (x, y)
370            region (None | int): alternative value for search region,
371                overwriting the `__init__` parameter `region`
372        """
373
374        if filtered is None or filtered.size == 0:
375            raise ValueError("Filtered image is empty or None")
376
377        region_size = region if region is not None else self._search_region
378        x, y = point
379
380        # Calculate crop boundaries
381        crop_x = max(0, x - region_size // 2)
382        crop_y = max(0, y - region_size // 2)
383        crop_width = min(region_size, filtered.shape[1] - crop_x)
384        crop_height = min(region_size, filtered.shape[0] - crop_y)
385
386        # Handle edge cases
387        if crop_width <= 0 or crop_height <= 0:
388            logger.warning(f"Point {point} is outside image bounds")
389            return point, 0.0
390
391        cropped = filtered[crop_y : crop_y + crop_height, crop_x : crop_x + crop_width]
392
393        if cropped.size == 0:
394            return point, 0.0
395
396        # Always apply Gaussian weighting by extending crop if needed
397        if cropped.shape[0] == region_size and cropped.shape[1] == region_size:
398            # Perfect size - apply weights directly
399            weights = self._create_gaussian_weights(region_size)
400            weighted = cropped.astype(np.float32) * weights
401        else:
402            # Extend crop to match region_size, apply weights, then restore
403            extended = np.zeros((region_size, region_size), dtype=cropped.dtype)
404
405            # Calculate offset to center the cropped region in extended array
406            offset_y = (region_size - cropped.shape[0]) // 2
407            offset_x = (region_size - cropped.shape[1]) // 2
408
409            # Place cropped region in center of extended array
410            extended[
411                offset_y : offset_y + cropped.shape[0],
412                offset_x : offset_x + cropped.shape[1],
413            ] = cropped
414
415            # Apply Gaussian weights to extended array
416            weights = self._create_gaussian_weights(region_size)
417            weighted_extended = extended.astype(np.float32) * weights
418
419            # Extract the original region back out
420            weighted = weighted_extended[
421                offset_y : offset_y + cropped.shape[0],
422                offset_x : offset_x + cropped.shape[1],
423            ]
424
425        best_idx = np.argmax(weighted)
426        best_y, best_x = np.unravel_index(best_idx, cropped.shape)
427
428        result_point = (
429            int(crop_x + best_x),
430            int(crop_y + best_y),
431        )
432        result_confidence = float(weighted[best_y, best_x]) / 255.0
433
434        return result_point, result_confidence
435
436    def find_table_points(
437        self,
438        img: MatLike | PathLike[str],
439        top_row: list[Point | None],
440        cell_widths: list[int],
441        cell_heights: list[int] | int,
442        visual: bool = False,
443        window: str = WINDOW,
444        goals_width: Optional[int] = None,
445        filtered: Optional[MatLike | PathLike[str]] = None,
446        smooth: bool = False,
447    ) -> "TableGrid":
448        """
449        Parse the image to a `TableGrid` structure that holds all of the
450        intersections between horizontal and vertical rules, starting near the `left_top` point
451
452        Args:
453            img (MatLike): the input image of a table
454            top_row (list[tuple[int, int]]): initial guess at topmost row of corner points (from template matching)
455            cell_widths (list[int]): the expected widths of the cells (based on a header template)
456            cell_heights (list[int]): the expected height of the rows of data.
457                The last value from this list is used until the image has no more vertical space.
458            visual (bool): whether to show intermediate steps
459            window (str): the name of the OpenCV window to use for visualization
460            goals_width (int | None): the width of the goal region when searching for the next point.
461                If None, defaults to 1.5 * search_region
462            filtered (MatLike | PathLike[str] | None): if provided, this image is used instead of
463                calculating the filtered image from scratch
464            smooth (bool): if True, smooth the grid after detection, using local heuristics
465
466        Returns:
467            a TableGrid object
468        """
469
470        if goals_width is None:
471            goals_width = self._search_region * 3 // 2
472
473        if not cell_widths:
474            raise ValueError("cell_widths must contain at least one value")
475
476        if not isinstance(img, np.ndarray):
477            img = cv.imread(os.fspath(img))
478
479        if filtered is None:
480            filtered = self.apply(img, visual)
481        else:
482            if not isinstance(filtered, np.ndarray):
483                filtered = cv.imread(os.fspath(filtered))
484
485            filtered = ensure_gray(filtered)
486
487        if visual:
488            imu.show(filtered, window=window)
489
490        if isinstance(cell_heights, int):
491            cell_heights = [cell_heights]
492
493        for i in range(len(top_row)):
494            if top_row[i] is None:
495                continue
496
497            adjusted, confidence = self.find_nearest(
498                filtered, top_row[i], int(self._search_region * 2)
499            )
500
501            if confidence < 0.15:
502                top_row[i] = None
503            else:
504                top_row[i] = adjusted
505
506        if not any(top_row):
507            logger.error("No good starting candidates given")
508
509        # resize all parameters according to scale
510        img = cv.resize(img, None, fx=self._scale, fy=self._scale)
511
512        if visual:
513            imu.push(img)
514
515        filtered = cv.resize(filtered, None, fx=self._scale, fy=self._scale)
516        cell_widths = [int(w * self._scale) for w in cell_widths]
517        cell_heights = [int(h * self._scale) for h in cell_heights]
518        top_row = [
519            (int(p[0] * self._scale), int(p[1] * self._scale))
520            if p is not None
521            else None
522            for p in top_row
523        ]
524        search_region = int(self._search_region * self._scale)
525
526        img_gray = ensure_gray(img)
527        filtered_gray = ensure_gray(filtered)
528
529        table_grower = TableGrower(
530            filtered_gray,
531            cell_widths,  # pyright: ignore
532            cell_heights,  # pyright: ignore
533            top_row,
534            search_region,
535            self._distance_penalty,
536            self._look_distance,
537            self._grow_threshold,
538            self._skip_astar_threshold,
539            self._min_rows,
540            self._cuts,
541            self._cut_fraction,
542        )
543
544        def show_grower_progress(wait: bool = False):
545            img_orig = np.copy(img)
546            corners = table_grower.get_all_corners()
547            for y in range(len(corners)):
548                for x in range(len(corners[y])):
549                    if corners[y][x] is not None:
550                        img_orig = imu.draw_points(
551                            img_orig,
552                            [corners[y][x]],
553                            color=(0, 0, 255),
554                            thickness=30,
555                        )
556
557            edge = table_grower.get_edge_points()
558
559            for point, score in edge:
560                color = (100, int(clamp(score * 255, 0, 255)), 100)
561                imu.draw_point(img_orig, point, color=color, thickness=20)
562
563            imu.show(img_orig, wait=wait)
564
565        if visual:
566            threshold = self._grow_threshold
567            look_distance = self._look_distance
568
569            # python implementation of rust loops, for visualization purposes
570            # note this is a LOT slower
571            while table_grower.grow_point(img_gray, filtered_gray) is not None:
572                show_grower_progress()
573
574            show_grower_progress(True)
575
576            original_threshold = threshold
577
578            loops_without_change = 0
579
580            while not table_grower.is_table_complete():
581                loops_without_change += 1
582
583                if loops_without_change > 50:
584                    break
585
586                if table_grower.extrapolate_one(img_gray, filtered_gray) is not None:
587                    show_grower_progress()
588
589                    loops_without_change = 0
590
591                    grown = False
592                    while table_grower.grow_point(img_gray, filtered_gray) is not None:
593                        show_grower_progress()
594                        grown = True
595                        threshold = min(0.1 + 0.9 * threshold, original_threshold)
596                        table_grower.set_threshold(threshold)
597
598                    if not grown:
599                        threshold *= 0.9
600                        table_grower.set_threshold(threshold)
601
602                else:
603                    threshold *= 0.9
604                    table_grower.set_threshold(threshold)
605
606                    if table_grower.grow_point(img_gray, filtered_gray) is not None:
607                        show_grower_progress()
608                        loops_without_change = 0
609
610        else:
611            table_grower.grow_table(img_gray, filtered_gray)
612
613        if smooth:
614            table_grower.smooth_grid()
615        corners = table_grower.get_all_corners()
616        logger.info(
617            f"Table growth complete, found {len(corners)} rows and {len(corners[0])} columns"
618        )
619        # rescale corners back to original size
620        if self._scale != 1.0:
621            for y in range(len(corners)):
622                for x in range(len(corners[y])):
623                    if corners[y][x] is not None:
624                        corners[y][x] = (
625                            int(corners[y][x][0] / self._scale),  # pyright:ignore
626                            int(corners[y][x][1] / self._scale),  # pyright:ignore
627                        )
628
629        return TableGrid(corners)  # pyright: ignore
630
631    def _visualize_grid(self, img: MatLike, points: List[List[Point]]) -> None:
632        """Visualize the detected grid points."""
633        all_points = [point for row in points for point in row]
634        drawn = imu.draw_points(img, all_points)
635        imu.show(drawn, wait=True)
636
637    def _visualize_path_finding(
638        self,
639        path: List[Point],
640        current: Point,
641        next_point: Point,
642        previous_row_target: Optional[Point] = None,
643        region_center: Optional[Point] = None,
644        region_size: Optional[int] = None,
645    ) -> None:
646        """Visualize the path finding process for debugging."""
647        global show_time
648
649        screen = imu.pop()
650
651        # if gray, convert to BGR
652        if len(screen.shape) == 2 or screen.shape[2] == 1:
653            debug_img = cv.cvtColor(screen, cv.COLOR_GRAY2BGR)
654        else:
655            debug_img = cast(MatLike, screen)
656
657        debug_img = imu.draw_points(debug_img, path, color=(200, 200, 0), thickness=2)
658        debug_img = imu.draw_points(
659            debug_img, [current], color=(0, 255, 0), thickness=3
660        )
661        debug_img = imu.draw_points(
662            debug_img, [next_point], color=(0, 0, 255), thickness=2
663        )
664
665        # Draw previous row target if available
666        if previous_row_target is not None:
667            debug_img = imu.draw_points(
668                debug_img, [previous_row_target], color=(255, 0, 255), thickness=2
669            )
670
671        # Draw search region if available
672        if region_center is not None and region_size is not None:
673            top_left = (
674                max(0, region_center[0] - region_size // 2),
675                max(0, region_center[1] - region_size // 2),
676            )
677            bottom_right = (
678                min(debug_img.shape[1], region_center[0] + region_size // 2),
679                min(debug_img.shape[0], region_center[1] + region_size // 2),
680            )
681            cv.rectangle(
682                debug_img,
683                top_left,
684                bottom_right,
685                color=(255, 0, 0),
686                thickness=2,
687                lineType=cv.LINE_AA,
688            )
689
690        imu.push(debug_img)
691
692        show_time += 1
693        if show_time % 10 != 1:
694            return
695
696        imu.show(debug_img, title="Next column point", wait=False)
697        # time.sleep(0.003)
698
699    @log_calls(level=logging.DEBUG, include_return=True)
700    def _astar(
701        self,
702        img: np.ndarray,
703        start: tuple[int, int],
704        goals: list[tuple[int, int]],
705        direction: str,
706    ) -> Optional[List[Point]]:
707        """
708        Find the best path between the start point and one of the goal points on the image
709        """
710
711        if not goals:
712            return None
713
714        if self._scale != 1.0:
715            img = cv.resize(img, None, fx=self._scale, fy=self._scale)
716            start = (int(start[0] * self._scale), int(start[1] * self._scale))
717            goals = [(int(g[0] * self._scale), int(g[1] * self._scale)) for g in goals]
718
719        # calculate bounding box with margin
720        all_points = goals + [start]
721        xs = [p[0] for p in all_points]
722        ys = [p[1] for p in all_points]
723
724        margin = 30
725        top_left = (max(0, min(xs) - margin), max(0, min(ys) - margin))
726        bottom_right = (
727            min(img.shape[1], max(xs) + margin),
728            min(img.shape[0], max(ys) + margin),
729        )
730
731        # check bounds
732        if (
733            top_left[0] >= bottom_right[0]
734            or top_left[1] >= bottom_right[1]
735            or top_left[0] >= img.shape[1]
736            or top_left[1] >= img.shape[0]
737        ):
738            return None
739
740        # transform coordinates to cropped image
741        start_local = (start[0] - top_left[0], start[1] - top_left[1])
742        goals_local = [(g[0] - top_left[0], g[1] - top_left[1]) for g in goals]
743
744        cropped = img[top_left[1] : bottom_right[1], top_left[0] : bottom_right[0]]
745
746        if cropped.size == 0:
747            return None
748
749        path = rust_astar(cropped, start_local, goals_local, direction)
750
751        if path is None:
752            return None
753
754        if self._scale != 1.0:
755            path = [(int(p[0] / self._scale), int(p[1] / self._scale)) for p in path]
756            top_left = (int(top_left[0] / self._scale), int(top_left[1] / self._scale))
757
758        return [(p[0] + top_left[0], p[1] + top_left[1]) for p in path]