Просмотр исходного кода

Split extract_crossword into extract_crossword_grid and grid_to_png
functions.

Andrew Klopper лет назад: 6
Родитель
Сommit
6fdfd76066
3 измененных файлов с 25 добавлено и 11 удалено
  1. 8 3
      home/management/commands/extract.py
  2. 3 2
      home/views.py
  3. 14 6
      home/xword.py

+ 8 - 3
home/management/commands/extract.py

1
 from django.core.management.base import BaseCommand
1
 from django.core.management.base import BaseCommand
2
-from ...xword import extract_crossword
2
+from ...xword import extract_crossword_grid, grid_to_png
3
 
3
 
4
 
4
 
5
 class Command(BaseCommand):
5
 class Command(BaseCommand):
37
         group.add_argument('--html')
37
         group.add_argument('--html')
38
 
38
 
39
     def handle(self, *args, **options):
39
     def handle(self, *args, **options):
40
-        image, warnings = extract_crossword(
40
+        warnings, grid, num_rows, num_cols, block_img = extract_crossword_grid(
41
             options['input_file_name'],
41
             options['input_file_name'],
42
             gaussian_blur_size=options['gaussian_blur_size'],
42
             gaussian_blur_size=options['gaussian_blur_size'],
43
             adaptive_threshold_block_size=options['adaptive_threshold_block_size'],
43
             adaptive_threshold_block_size=options['adaptive_threshold_block_size'],
49
             line_detector_element_size=options['line_detector_element_size'],
49
             line_detector_element_size=options['line_detector_element_size'],
50
             sampling_block_size_ratio=options['sampling_block_size_ratio'],
50
             sampling_block_size_ratio=options['sampling_block_size_ratio'],
51
             sampling_threshold_quantile=options['sampling_threshold_quantile'],
51
             sampling_threshold_quantile=options['sampling_threshold_quantile'],
52
-            sampling_threshold=options['sampling_threshold'],
52
+            sampling_threshold=options['sampling_threshold']
53
+        )
54
+        image = grid_to_png(
55
+            grid,
56
+            num_rows,
57
+            num_cols,
53
             grid_line_thickness=options['grid_line_thickness'],
58
             grid_line_thickness=options['grid_line_thickness'],
54
             grid_square_size=options['grid_square_size'],
59
             grid_square_size=options['grid_square_size'],
55
             grid_border_size=options['grid_border_size']
60
             grid_border_size=options['grid_border_size']

+ 3 - 2
home/views.py

4
 from django.shortcuts import render
4
 from django.shortcuts import render
5
 from django.views import View
5
 from django.views import View
6
 from .forms import CrosswordForm
6
 from .forms import CrosswordForm
7
-from .xword import extract_crossword
7
+from .xword import extract_crossword_grid, grid_to_png
8
 
8
 
9
 
9
 
10
 class HomeView(View):
10
 class HomeView(View):
15
         form = CrosswordForm(request.POST, request.FILES)
15
         form = CrosswordForm(request.POST, request.FILES)
16
         if not form.is_valid():
16
         if not form.is_valid():
17
             return HttpResponseBadRequest('Invalid form data')
17
             return HttpResponseBadRequest('Invalid form data')
18
-        image, warnings = extract_crossword(
18
+        warnings, grid, num_rows, num_cols, block_img = extract_crossword_grid(
19
             form.cleaned_data['file'].temporary_file_path(),
19
             form.cleaned_data['file'].temporary_file_path(),
20
             filter_colours=False
20
             filter_colours=False
21
         )
21
         )
22
+        image = grid_to_png(grid, num_rows, num_cols)
22
         return render(request, 'home/output.html', {
23
         return render(request, 'home/output.html', {
23
             'warnings': warnings,
24
             'warnings': warnings,
24
             'image_file_name': 'xword_{}.png'.format(datetime.now().strftime('%Y%m%d_%H%M%S')),
25
             'image_file_name': 'xword_{}.png'.format(datetime.now().strftime('%Y%m%d_%H%M%S')),

+ 14 - 6
home/xword.py

196
     cv2.destroyAllWindows()
196
     cv2.destroyAllWindows()
197
 
197
 
198
 
198
 
199
-def extract_crossword(
199
+def extract_crossword_grid(
200
     file_name,
200
     file_name,
201
     filter_colours=False,
201
     filter_colours=False,
202
     colour_filter_threshold=48,
202
     colour_filter_threshold=48,
210
     line_detector_element_size=51,
210
     line_detector_element_size=51,
211
     sampling_block_size_ratio=0.25,
211
     sampling_block_size_ratio=0.25,
212
     sampling_threshold_quantile=0.3,
212
     sampling_threshold_quantile=0.3,
213
-    sampling_threshold=None,
214
-    grid_line_thickness=4,
215
-    grid_square_size=64,
216
-    grid_border_size=20,
213
+    sampling_threshold=None
217
 ):
214
 ):
218
     warnings = []
215
     warnings = []
219
 
216
 
245
     if warning:
242
     if warning:
246
         warnings.append("Some blocks may be the wrong colour")
243
         warnings.append("Some blocks may be the wrong colour")
247
 
244
 
245
+    return warnings, grid, num_rows, num_cols, block_img
246
+
247
+
248
+def grid_to_png(
249
+    grid,
250
+    num_rows,
251
+    num_cols,
252
+    grid_line_thickness=4,
253
+    grid_square_size=64,
254
+    grid_border_size=20
255
+):
248
     step = grid_square_size + grid_line_thickness
256
     step = grid_square_size + grid_line_thickness
249
     grid_height = num_rows * step + grid_line_thickness
257
     grid_height = num_rows * step + grid_line_thickness
250
     grid_width = num_cols * step + grid_line_thickness
258
     grid_width = num_cols * step + grid_line_thickness
258
                 cv2.rectangle(output, (x, y), (x + grid_square_size - 1, y + grid_square_size - 1), 255, -1)
266
                 cv2.rectangle(output, (x, y), (x + grid_square_size - 1, y + grid_square_size - 1), 255, -1)
259
 
267
 
260
     _, png = cv2.imencode('.png', output)
268
     _, png = cv2.imencode('.png', output)
261
-    return png.tobytes(), warnings
269
+    return png.tobytes()