| import unittest |
| from dataclasses import dataclass, field |
|
|
| from ocr import drop_short_lines |
|
|
|
|
| @dataclass |
| class Line: |
| baseline: list |
|
|
|
|
| @dataclass |
| class Bounds: |
| lines: list = field(default_factory=list) |
|
|
|
|
| def row(x0, x1, y=100): |
| return Line(baseline=[[x0, y], [x1, y]]) |
|
|
|
|
| class DropShortLinesTests(unittest.TestCase): |
| def test_drops_marginal_fragments(self): |
| bounds = Bounds([row(100, 1100, y) for y in range(100, 1300, 100)] + [row(20, 60), row(1150, 1190)]) |
| kept, dropped = drop_short_lines(bounds, 0.2) |
| self.assertEqual(dropped, 2) |
| self.assertEqual(len(kept.lines), 12) |
|
|
| def test_keeps_short_final_line_above_threshold(self): |
| bounds = Bounds([row(100, 1100, y) for y in range(100, 900, 100)] + [row(100, 400, 1000)]) |
| kept, dropped = drop_short_lines(bounds, 0.2) |
| self.assertEqual(dropped, 0) |
| self.assertEqual(len(kept.lines), 9) |
|
|
| def test_disabled_or_empty_is_unchanged(self): |
| bounds = Bounds([row(0, 10), row(0, 1000)]) |
| self.assertEqual(drop_short_lines(bounds, 0)[1], 0) |
| self.assertEqual(drop_short_lines(Bounds([]), 0.2)[1], 0) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
