Quellcode durchsuchen

added --split-auto

master
Evgeniy Ierusalimov vor 1 Woche
Ursprung
Commit
319c26cb0d
4 geänderte Dateien mit 153 neuen und 6 gelöschten Zeilen
  1. 31
    3
      src/cli.py
  2. 2
    0
      src/split/__init__.py
  3. 55
    3
      src/split/slicer.py
  4. 65
    0
      tests/test_slicer.py

+ 31
- 3
src/cli.py Datei anzeigen

8
 from src.split.slicer import (
8
 from src.split.slicer import (
9
     DEFAULT_BORDER_PX,
9
     DEFAULT_BORDER_PX,
10
     VALID_ROTATIONS,
10
     VALID_ROTATIONS,
11
+    detect_grid,
12
+    load_image,
11
     parse_slice,
13
     parse_slice,
12
     process_image,
14
     process_image,
15
+    rotate_image,
13
 )
16
 )
14
 
17
 
15
 app = typer.Typer(add_completion=False)
18
 app = typer.Typer(add_completion=False)
30
         ),
33
         ),
31
     ],
34
     ],
32
     slice: Annotated[
35
     slice: Annotated[
33
-        str,
36
+        str | None,
34
         typer.Option(
37
         typer.Option(
35
             "--slice",
38
             "--slice",
36
             "-s",
39
             "-s",
37
             help="Размер сетки в формате <колонки>:<строки>, например 3:2",
40
             help="Размер сетки в формате <колонки>:<строки>, например 3:2",
38
         ),
41
         ),
39
-    ],
42
+    ] = None,
43
+    slice_auto: Annotated[
44
+        bool,
45
+        typer.Option(
46
+            "--slice-auto",
47
+            "-a",
48
+            help="Автоматически определить размер сетки по содержимому изображения",
49
+        ),
50
+    ] = False,
40
     output_dir: Annotated[
51
     output_dir: Annotated[
41
         Path,
52
         Path,
42
         typer.Option(
53
         typer.Option(
80
             f"Недопустимый угол: {pre_rotate}. Допустимые: {sorted(VALID_ROTATIONS)}"
91
             f"Недопустимый угол: {pre_rotate}. Допустимые: {sorted(VALID_ROTATIONS)}"
81
         )
92
         )
82
 
93
 
83
-    cols, rows = parse_slice(slice)
94
+    if slice is not None and slice_auto:
95
+        raise typer.BadParameter(
96
+            "Нельзя указывать одновременно --slice и --slice-auto"
97
+        )
98
+    if slice is None and not slice_auto:
99
+        raise typer.BadParameter(
100
+            "Укажите --slice или --slice-auto"
101
+        )
102
+
103
+    if slice_auto:
104
+        image = load_image(input.resolve())
105
+        if pre_rotate is not None:
106
+            image = rotate_image(image, pre_rotate)
107
+        cols, rows = detect_grid(image)
108
+        typer.echo(f"Определена сетка: {cols}×{rows}")
109
+    else:
110
+        assert slice is not None
111
+        cols, rows = parse_slice(slice)
84
 
112
 
85
     output_paths = process_image(
113
     output_paths = process_image(
86
         input_path=input.resolve(),
114
         input_path=input.resolve(),

+ 2
- 0
src/split/__init__.py Datei anzeigen

6
     VALID_ROTATIONS,
6
     VALID_ROTATIONS,
7
     add_border,
7
     add_border,
8
     crop_to_content,
8
     crop_to_content,
9
+    detect_grid,
9
     find_content_bounds,
10
     find_content_bounds,
10
     generate_output_paths,
11
     generate_output_paths,
11
     load_image,
12
     load_image,
25
     "VALID_ROTATIONS",
26
     "VALID_ROTATIONS",
26
     "add_border",
27
     "add_border",
27
     "crop_to_content",
28
     "crop_to_content",
29
+    "detect_grid",
28
     "find_content_bounds",
30
     "find_content_bounds",
29
     "generate_output_paths",
31
     "generate_output_paths",
30
     "load_image",
32
     "load_image",

+ 55
- 3
src/split/slicer.py Datei anzeigen

12
 VALID_ROTATIONS: frozenset[int] = frozenset({90, 180, 270})
12
 VALID_ROTATIONS: frozenset[int] = frozenset({90, 180, 270})
13
 CONTENT_THRESHOLD: int = 200
13
 CONTENT_THRESHOLD: int = 200
14
 DENSITY_RATIO: float = 0.005
14
 DENSITY_RATIO: float = 0.005
15
+AUTO_GRID_MAX: int = 4
15
 
16
 
16
 
17
 
17
 def load_image(path: Path) -> np.ndarray:
18
 def load_image(path: Path) -> np.ndarray:
115
     return path
116
     return path
116
 
117
 
117
 
118
 
119
+def _binarize(image: np.ndarray) -> np.ndarray:
120
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
121
+    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
122
+    otsu_th, binary = cv2.threshold(
123
+        blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU,
124
+    )
125
+    if otsu_th < 50:
126
+        binary = cv2.adaptiveThreshold(
127
+            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
128
+            cv2.THRESH_BINARY_INV, 31, 10,
129
+        )
130
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
131
+    return cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
132
+
133
+
134
+def detect_grid(image: np.ndarray) -> tuple[int, int]:
135
+    binary = _binarize(image)
136
+
137
+    cols = _count_bands(binary, axis=0)
138
+    rows = _count_bands(binary, axis=1)
139
+
140
+    return cols, rows
141
+
142
+
143
+def _count_bands(binary: np.ndarray, axis: int) -> int:
144
+    other_dim = binary.shape[axis]
145
+    length = binary.shape[1 - axis]
146
+    projection = (binary > 0).sum(axis=axis).astype(np.float64)
147
+    density = projection / other_dim
148
+
149
+    dilate_width = max(5, length // 80)
150
+    kernel_1d = np.ones(dilate_width, dtype=np.uint8)
151
+    density_u8 = (density * 255).clip(0, 255).astype(np.uint8)
152
+    density_2d = density_u8.reshape(1, -1)
153
+    dilated = cv2.dilate(density_2d, kernel_1d)
154
+    is_content = (dilated[0] > 1).tolist()
155
+
156
+    bands = 0
157
+    in_band = False
158
+    for val in is_content:
159
+        if val and not in_band:
160
+            bands += 1
161
+            in_band = True
162
+        elif not val:
163
+            in_band = False
164
+    return max(1, bands)
165
+
166
+
118
 def find_content_bounds(
167
 def find_content_bounds(
119
     image: np.ndarray,
168
     image: np.ndarray,
120
     threshold: int = CONTENT_THRESHOLD,
169
     threshold: int = CONTENT_THRESHOLD,
175
 
224
 
176
 def process_image(
225
 def process_image(
177
     input_path: Path,
226
     input_path: Path,
178
-    rows: int,
179
-    cols: int,
180
-    output_dir: Path,
227
+    rows: int | None = None,
228
+    cols: int | None = None,
229
+    output_dir: Path = Path(),
181
     pre_rotate: int | None = None,
230
     pre_rotate: int | None = None,
182
     border_px: int = DEFAULT_BORDER_PX,
231
     border_px: int = DEFAULT_BORDER_PX,
183
     post_crop: bool = False,
232
     post_crop: bool = False,
187
     if pre_rotate is not None:
236
     if pre_rotate is not None:
188
         image = rotate_image(image, pre_rotate)
237
         image = rotate_image(image, pre_rotate)
189
 
238
 
239
+    if rows is None or cols is None:
240
+        cols, rows = detect_grid(image)
241
+
190
     pages = slice_grid(image, rows, cols)
242
     pages = slice_grid(image, rows, cols)
191
 
243
 
192
     if post_crop:
244
     if post_crop:

+ 65
- 0
tests/test_slicer.py Datei anzeigen

12
     VALID_ROTATIONS,
12
     VALID_ROTATIONS,
13
     add_border,
13
     add_border,
14
     crop_to_content,
14
     crop_to_content,
15
+    detect_grid,
15
     find_content_bounds,
16
     find_content_bounds,
16
     generate_output_paths,
17
     generate_output_paths,
17
     load_image,
18
     load_image,
483
                 n = cv2.imread(str(np))
484
                 n = cv2.imread(str(np))
484
                 assert c.shape[0] < n.shape[0]
485
                 assert c.shape[0] < n.shape[0]
485
                 assert c.shape[1] < n.shape[1]
486
                 assert c.shape[1] < n.shape[1]
487
+
488
+
489
+class TestDetectGrid:
490
+    def test_2x2(self) -> None:
491
+        image = _create_test_image(400, 300)
492
+        image[:, :] = (255, 255, 255)
493
+        dark = (0, 0, 0)
494
+        gap = 20
495
+        bw = (400 - 3 * gap) // 2
496
+        bh = (300 - 3 * gap) // 2
497
+        for r in range(2):
498
+            for c in range(2):
499
+                y = gap + r * (bh + gap)
500
+                x = gap + c * (bw + gap)
501
+                image[y:y + bh, x:x + bw] = dark
502
+        cols, rows = detect_grid(image)
503
+        assert cols == 2
504
+        assert rows == 2
505
+
506
+    def test_2x4(self) -> None:
507
+        image = _create_test_image(800, 400)
508
+        image[:, :] = (255, 255, 255)
509
+        gap = 15
510
+        bw = (800 - 5 * gap) // 4
511
+        bh = (400 - 3 * gap) // 2
512
+        dark = (0, 0, 0)
513
+        for r in range(2):
514
+            for c in range(4):
515
+                y = gap + r * (bh + gap)
516
+                x = gap + c * (bw + gap)
517
+                image[y:y + bh, x:x + bw] = dark
518
+        cols, rows = detect_grid(image)
519
+        assert cols == 4
520
+        assert rows == 2
521
+
522
+    def test_1x1(self) -> None:
523
+        image = _create_test_image(200, 200)
524
+        image[:, :] = (255, 255, 255)
525
+        image[20:180, 20:180] = (0, 0, 0)
526
+        cols, rows = detect_grid(image)
527
+        assert cols == 1
528
+        assert rows == 1
529
+
530
+    def test_3x3(self) -> None:
531
+        image = _create_test_image(600, 450)
532
+        image[:, :] = (255, 255, 255)
533
+        dark = (0, 0, 0)
534
+        gap = 10
535
+        bw = (600 - 4 * gap) // 3
536
+        bh = (450 - 4 * gap) // 3
537
+        for r in range(3):
538
+            for c in range(3):
539
+                y = gap + r * (bh + gap)
540
+                x = gap + c * (bw + gap)
541
+                image[y:y + bh, x:x + bw] = dark
542
+        cols, rows = detect_grid(image)
543
+        assert cols == 3
544
+        assert rows == 3
545
+
546
+    def test_uniform_image(self) -> None:
547
+        image = _create_test_image()
548
+        cols, rows = detect_grid(image)
549
+        assert cols == 1
550
+        assert rows == 1

Laden…
Abbrechen
Speichern