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