xword/home/management/commands/extract.py

from django.core.management.base import BaseCommand
from ...xword import extract_crossword_grid, grid_to_png


class Command(BaseCommand):
    help = "Extracts a clean crossword image from a photograph."

    def add_arguments(self, parser):
        parser.add_argument('input_file_name')
        parser.add_argument('output_file_name')

        parser.add_argument('--filter-colours', action='store_true')
        parser.add_argument('--colour-filter-threshold', type=int, default=48)

        parser.add_argument('--gaussian-blur-size', type=int, default=11)
        parser.add_argument('--adaptive-threshold-block-size', type=int, default=11)
        parser.add_argument('--adaptive-threshold-mean-adjustment', type=int, default=2)

        parser.add_argument('--not-square', action='store_true')

        parser.add_argument('--num-dilations', type=int, default=1)

        parser.add_argument('--contour-erosion-kernel-size', type=int, default=5)
        parser.add_argument('--contour-erosion-iterations', type=int, default=5)
        parser.add_argument('--line-detector-element-size', type=int, default=51)

        parser.add_argument('--sampling-block-size-ratio', type=float, default=0.25)
        parser.add_argument('--sampling-threshold-quantile', type=float, default=0.3)
        parser.add_argument('--sampling-threshold', type=int)

        parser.add_argument('--grid-line-thickness', type=int, default=4)
        parser.add_argument('--grid-square-size', type=int, default=64)
        parser.add_argument('--grid-border-size', type=int, default=20)

        group = parser.add_mutually_exclusive_group()
        group.add_argument('--out')
        group.add_argument('--html')

    def handle(self, *args, **options):
        warnings, grid, num_rows, num_cols, block_img = extract_crossword_grid(
            options['input_file_name'],
            gaussian_blur_size=options['gaussian_blur_size'],
            adaptive_threshold_block_size=options['adaptive_threshold_block_size'],
            adaptive_threshold_mean_adjustment=options['adaptive_threshold_mean_adjustment'],
            square=not options['not_square'],
            num_dilations=options['num_dilations'],
            contour_erosion_kernel_size=options['contour_erosion_kernel_size'],
            contour_erosion_iterations=options['contour_erosion_iterations'],
            line_detector_element_size=options['line_detector_element_size'],
            sampling_block_size_ratio=options['sampling_block_size_ratio'],
            sampling_threshold_quantile=options['sampling_threshold_quantile'],
            sampling_threshold=options['sampling_threshold']
        )
        image = grid_to_png(
            grid,
            num_rows,
            num_cols,
            grid_line_thickness=options['grid_line_thickness'],
            grid_square_size=options['grid_square_size'],
            grid_border_size=options['grid_border_size']
        )
        for warning in warnings:
            print('WARNING: ' + warning)
        with open(options['output_file_name'], 'wb') as f:
            f.write(image)