Coverage for collectra/tasks/line_detection.py: 95.12%
164 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Deterministic printed-text line detection."""
3from __future__ import annotations
5from typing import Any
7import numpy as np
9from collectra.tasks.base import Task
10from collectra.types.images import Image, ImageCrop, Orientation
12__all__ = ["LineDetectorRLSA"]
15class LineDetectorRLSA(Task):
16 """Split printed text into lines using smearing and projection profiles.
18 The detector is deterministic and does not load a machine-learning model.
19 All size parameters are relative to the input dimensions so that the same
20 configuration works across scan resolutions.
21 """
23 def __init__(
24 self,
25 name: str,
26 deskew: bool = True,
27 max_skew_degrees: float = 5.0,
28 adaptive_block_fraction: float = 0.03,
29 adaptive_constant: float = 15.0,
30 rule_length_fraction: float = 0.35,
31 smear_width_fraction: float = 0.015,
32 smear_height_fraction: float = 0.003,
33 smooth_height_fraction: float = 0.002,
34 min_ink_fraction: float = 0.002,
35 min_component_area_fraction: float = 0.000002,
36 min_line_height_fraction: float = 0.003,
37 max_line_gap_fraction: float = 0.004,
38 padding_fraction: float = 0.003,
39 **kwargs: Any,
40 ) -> None:
41 super().__init__(name, **kwargs)
42 self.deskew = deskew
43 self.max_skew_degrees = max_skew_degrees
44 self.adaptive_block_fraction = adaptive_block_fraction
45 self.adaptive_constant = adaptive_constant
46 self.rule_length_fraction = rule_length_fraction
47 self.smear_width_fraction = smear_width_fraction
48 self.smear_height_fraction = smear_height_fraction
49 self.smooth_height_fraction = smooth_height_fraction
50 self.min_ink_fraction = min_ink_fraction
51 self.min_component_area_fraction = min_component_area_fraction
52 self.min_line_height_fraction = min_line_height_fraction
53 self.max_line_gap_fraction = max_line_gap_fraction
54 self.padding_fraction = padding_fraction
55 self._validate_parameters()
57 def _validate_parameters(self) -> None:
58 fractions = (
59 "adaptive_block_fraction",
60 "rule_length_fraction",
61 "smear_width_fraction",
62 "smear_height_fraction",
63 "smooth_height_fraction",
64 "min_ink_fraction",
65 "min_component_area_fraction",
66 "min_line_height_fraction",
67 "max_line_gap_fraction",
68 "padding_fraction",
69 )
70 for name in fractions:
71 value = getattr(self, name)
72 if not 0 <= value <= 1:
73 raise ValueError(f"{name} must be between 0 and 1, got {value}")
74 if not 0 <= self.max_skew_degrees <= 45:
75 raise ValueError("max_skew_degrees must be between 0 and 45")
77 def input_type(self) -> type:
78 return Image
80 def output_type(self) -> type:
81 return ImageCrop
83 @staticmethod
84 def _odd_size(value: float, minimum: int = 3) -> int:
85 size = max(minimum, int(round(value)))
86 return size if size % 2 else size + 1
88 def _binarize(self, image: np.ndarray) -> np.ndarray:
89 import cv2
91 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
92 block_size = self._odd_size(
93 min(gray.shape) * self.adaptive_block_fraction, minimum=15
94 )
95 binary = cv2.adaptiveThreshold(
96 gray,
97 255,
98 cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
99 cv2.THRESH_BINARY_INV,
100 block_size,
101 self.adaptive_constant,
102 )
103 return self._remove_noise(binary)
105 def _remove_noise(self, binary: np.ndarray) -> np.ndarray:
106 import cv2
108 height, width = binary.shape
109 minimum_area = max(2, round(height * width * self.min_component_area_fraction))
110 count, labels, stats, _ = cv2.connectedComponentsWithStats(binary, 8)
111 cleaned = np.zeros_like(binary)
112 for label in range(1, count):
113 if stats[label, cv2.CC_STAT_AREA] >= minimum_area:
114 cleaned[labels == label] = 255
115 return cleaned
117 def _remove_horizontal_rules(self, binary: np.ndarray) -> np.ndarray:
118 import cv2
120 width = binary.shape[1]
121 kernel_width = max(3, round(width * self.rule_length_fraction))
122 kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_width, 1))
123 rules = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
124 return cv2.subtract(binary, rules)
126 def _estimate_skew(self, binary: np.ndarray) -> float:
127 import cv2
129 if not self.deskew or not np.any(binary):
130 return 0.0
131 height, width = binary.shape
132 lines = cv2.HoughLinesP(
133 binary,
134 1,
135 np.pi / 720,
136 threshold=max(20, width // 20),
137 minLineLength=max(20, width // 12),
138 maxLineGap=max(5, width // 100),
139 )
140 if lines is None:
141 return 0.0
142 angles = []
143 # reshape, not [:, 0] — cv2.HoughLinesP returns (N, 1, 4) on OpenCV 4.x
144 # but (N, 4) on 5.x; reshape(-1, 4) normalises either into one row per line.
145 for x1, y1, x2, y2 in lines.reshape(-1, 4):
146 angle = float(np.degrees(np.arctan2(y2 - y1, x2 - x1)))
147 if abs(angle) <= self.max_skew_degrees:
148 angles.append(angle)
149 return float(np.median(angles)) if angles else 0.0
151 @staticmethod
152 def _rotate(binary: np.ndarray, angle: float) -> np.ndarray:
153 import cv2
155 if abs(angle) < 0.01:
156 return binary
157 height, width = binary.shape
158 matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
159 return cv2.warpAffine(
160 binary,
161 matrix,
162 (width, height),
163 flags=cv2.INTER_NEAREST,
164 borderMode=cv2.BORDER_CONSTANT,
165 borderValue=0,
166 )
168 def _line_bands(self, binary: np.ndarray) -> list[tuple[int, int]]:
169 import cv2
171 height, width = binary.shape
172 kernel = cv2.getStructuringElement(
173 cv2.MORPH_RECT,
174 (
175 max(1, round(width * self.smear_width_fraction)),
176 max(1, round(height * self.smear_height_fraction)),
177 ),
178 )
179 smeared = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
180 profile = np.count_nonzero(smeared, axis=1).astype(float)
181 smooth_size = self._odd_size(height * self.smooth_height_fraction)
182 profile = np.convolve(profile, np.ones(smooth_size) / smooth_size, mode="same")
183 occupied = profile >= max(1, width * self.min_ink_fraction)
185 ranges: list[tuple[int, int]] = []
186 starts = np.flatnonzero(occupied & ~np.r_[False, occupied[:-1]])
187 ends = np.flatnonzero(occupied & ~np.r_[occupied[1:], False]) + 1
188 min_height = max(1, round(height * self.min_line_height_fraction))
189 for top, bottom in zip(starts, ends):
190 if bottom - top >= min_height:
191 ranges.append((int(top), int(bottom)))
193 max_gap = round(height * self.max_line_gap_fraction)
194 merged: list[tuple[int, int]] = []
195 for top, bottom in ranges:
196 if merged and top - merged[-1][1] <= max_gap:
197 merged[-1] = (merged[-1][0], bottom)
198 else:
199 merged.append((top, bottom))
200 return merged
202 def _crop_bounds(
203 self, bands: list[tuple[int, int]], width: int, height: int, skew: float
204 ) -> list[tuple[int, int, int, int]]:
205 if not bands:
206 return []
207 skew_padding = abs(np.tan(np.radians(skew))) * width / 2
208 padding = round(height * self.padding_fraction + skew_padding)
209 boundaries = [0]
210 boundaries.extend(
211 (left[1] + right[0]) // 2 for left, right in zip(bands, bands[1:])
212 )
213 boundaries.append(height)
214 return [
215 (
216 0,
217 max(0, min(top - padding, band_top - padding)),
218 width,
219 min(height, max(bottom + padding, band_bottom + padding)),
220 )
221 for (band_top, band_bottom), top, bottom in zip(
222 bands, boundaries, boundaries[1:]
223 )
224 ]
226 def _make_crop(
227 self,
228 image: Image,
229 bounds: tuple[int, int, int, int],
230 display_width: int,
231 display_height: int,
232 ) -> ImageCrop:
233 left, top, right, bottom = bounds
234 local_x = (left + right) / (2 * display_width)
235 local_y = (top + bottom) / (2 * display_height)
236 local_width = (right - left) / display_width
237 local_height = (bottom - top) / display_height
239 display_left = local_x - local_width / 2
240 display_top = local_y - local_height / 2
241 display_right = local_x + local_width / 2
242 display_bottom = local_y + local_height / 2
243 display_corners = (
244 (display_left, display_top),
245 (display_right, display_top),
246 (display_right, display_bottom),
247 (display_left, display_bottom),
248 )
250 def to_source(x: float, y: float) -> tuple[float, float]:
251 match image.orientation:
252 case Orientation.WEST:
253 return y, 1 - x
254 case Orientation.SOUTH:
255 return 1 - x, 1 - y
256 case Orientation.EAST:
257 return 1 - y, x
258 case _:
259 return x, y
261 source_corners = [to_source(x, y) for x, y in display_corners]
262 source_left = min(point[0] for point in source_corners)
263 source_top = min(point[1] for point in source_corners)
264 source_right = max(point[0] for point in source_corners)
265 source_bottom = max(point[1] for point in source_corners)
266 source_x = (source_left + source_right) / 2
267 source_y = (source_top + source_bottom) / 2
268 source_width = source_right - source_left
269 source_height = source_bottom - source_top
271 if isinstance(image, ImageCrop):
272 return ImageCrop(
273 name=self.get_output_name(),
274 data=image.data,
275 x_center=image.x_center + (source_x - 0.5) * image.width_relative,
276 y_center=image.y_center + (source_y - 0.5) * image.height_relative,
277 width_relative=source_width * image.width_relative,
278 height_relative=source_height * image.height_relative,
279 orientation=image.orientation,
280 )
281 return image.make_crop(
282 source_x,
283 source_y,
284 source_width,
285 source_height,
286 name=self.get_output_name(),
287 orientation=image.orientation,
288 )
290 def run(self, *images: Image) -> list[ImageCrop]:
291 """Return one axis-aligned crop for each detected printed-text line."""
292 results: list[ImageCrop] = []
293 for image in images:
294 if not isinstance(image, Image):
295 raise TypeError(f"image must be Image, got {type(image)}")
296 array = np.asarray(image.pil().convert("RGB"))
297 binary = self._binarize(array)
298 skew = self._estimate_skew(binary)
299 deskewed = self._rotate(binary, -skew)
300 bands = self._line_bands(self._remove_horizontal_rules(deskewed))
301 bounds = self._crop_bounds(bands, array.shape[1], array.shape[0], skew)
302 if not bounds:
303 bounds = [(0, 0, array.shape[1], array.shape[0])]
304 results.extend(
305 self._make_crop(image, bound, array.shape[1], array.shape[0])
306 for bound in bounds
307 )
308 return results