xword/home/management/commands/extract.py

import cv2
from django.core.management.base import BaseCommand
from django.template.loader import render_to_string
from ...xword import extract_crossword_grid, draw_grid


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('--debug', action='store_true')

        parser.add_argument('--remove-colours', action='store_true')
        parser.add_argument('--colour-removal-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'],
            callback=debug_callback if options['debug'] else None,
            remove_colours=options['remove_colours'],
            colour_removal_threshold=options['colour_removal_threshold'],
            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']
        )

        if options['debug']:
            debug_callback('square', block_img)

        for warning in warnings:
            print('WARNING: ' + warning)

        if options['html'] is not None:
            html = render_to_string('home/grid.html', {
                'grid': grid,
                'num_rows': num_rows,
                'num_cols': num_cols,
            })

            with open(options['html'], 'w') as f:
                f.write(html)

        else:
            image = draw_grid(
                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']
            )

            if options['debug'] or (options['out'] is None):
                debug_callback('output', image)

            if options['out'] is not None:
                _, png = cv2.imencode('.png', image)

                with open(options['out'], 'wb') as f:
                    f.write(png.tobytes())

        if options['debug'] or (options['out'] is None):
            while cv2.waitKey() & 0xFF != ord('q'):
                pass
            cv2.destroyAllWindows()


def debug_callback(name, image):
    cv2.namedWindow(name, cv2.WINDOW_NORMAL)
    cv2.imshow(name, image)