proxyscripts/border_detection.py

import dataclasses
from typing import Dict, List, Tuple

from PIL import ImageDraw, Image

from base_proxy import *


TOP = 'top'
BOTTOM = 'bottom'
LEFT = 'left'
RIGHT = 'right'


@dataclasses.dataclass
class PixelColorBuckets:
    blacks: int = 0
    whites: int = 0
    others: int = 0


def border_color_to_rgb(border_color: str) -> Tuple[int, int, int]:
    return {
        Borders.SOLID_BLACK: (0, 0, 0),
        Borders.SOLID_WHITE: (255, 255, 255),
    }.get(border_color)


@dataclasses.dataclass
class Borders:
    SOLID_BLACK = 'solid_black'
    SOLID_WHITE = 'solid_white'
    BORDERLESS = 'borderless'
    UNKNOWN = 'unknown'

    dpi: float = None
    color: Dict[str, str] = dataclasses.field(default_factory=dict)
    thickness: Dict[str, int] = dataclasses.field(default_factory=dict)

    def inform_color(self, edge: str, all_color_buckets: Dict[int, PixelColorBuckets]) -> None:
        any_others = any(b.others > 0 for b in all_color_buckets.values())
        any_blacks = any(b.blacks > 0 for b in all_color_buckets.values())
        any_whites = any(b.whites > 0 for b in all_color_buckets.values())
        if any_others:
            self.color[edge] = Borders.BORDERLESS
        elif any_blacks and any_whites:
            self.color[edge] = Borders.UNKNOWN
        elif any_blacks:
            self.color[edge] = Borders.SOLID_BLACK
        elif any_whites:
            self.color[edge] = Borders.SOLID_WHITE
        else:
            self.color[edge] = Borders.UNKNOWN

    def inform_depth(self, edge: str, depths: Dict[int, int]) -> None:
        self.thickness[edge] = min(depths.values())

    def get_border_thickness_in(self, edge: str) -> float:
        return self.thickness[edge] / self.dpi


@dataclasses.dataclass
class CardBox:
    left: int
    top: int
    width: int
    height: int

    def get_sample_xs(self, sample_count: int) -> List[int]:
        xs = []
        for i in range(sample_count):
            x = self.left + (i * (self.width - 1)) // (sample_count - 1)
            xs.append(x)
        return xs

    def get_sample_ys(self, sample_count: int) -> List[int]:
        ys = []
        for i in range(sample_count):
            y = self.top + (i * (self.height - 1)) // (sample_count - 1)
            ys.append(y)
        return ys


def categorize_pixels_in_line(im: Image, start: Tuple[int, int], increment: Tuple[int, int], depth: int) -> PixelColorBuckets:
    color_buckets = PixelColorBuckets()

    xy = start
    for n in range(depth):
        pixel = im.getpixel(xy)
        if pixel == (0, 0, 0):
            color_buckets.blacks += 1
        elif pixel == (255, 255, 255):
            color_buckets.whites += 1
        else:
            color_buckets.others += 1
        xy = (xy[0] + increment[0], xy[1] + increment[1])

    return color_buckets


def detect_depth_in_column(im: Image, color: str, fixed_x: int, start: int, increment: int, max_depth: int) -> int:
    expected_rgb = border_color_to_rgb(color)
    y = start
    for n in range(max_depth):
        pixel = im.getpixel((fixed_x, y))
        if expected_rgb != pixel:
            return n - 1
        y += increment
    else:
        return max_depth


def detect_borders(loaded_input_card: 'InputCardImage') -> Borders:
    dpi: int = loaded_input_card.input_dpi
    im: Image = loaded_input_card.original_im
    assert dpi is not None
    assert im is not None

    input_w = im.size[0]
    input_h = im.size[1]
    card_w = in_to_px(CARD_W_INCHES, dpi)
    card_h = in_to_px(CARD_H_INCHES, dpi)
    bleed_left = (input_w - card_w) // 2
    bleed_top = (input_h - card_h) // 2
    card_box = CardBox(bleed_left, bleed_top, card_w, card_h)

    if SAVE_INTERMEDIATE_FOR_DEBUG:
        draw = ImageDraw.Draw(im)

    borders = Borders()
    borders.dpi = dpi

    depth_to_calibrate = mil_to_px(10, dpi)
    max_depth_to_check = in_to_px(0.5, dpi)
    border_depth = 0

    # Sample 30 lines
    xs = card_box.get_sample_xs(30)
    ys = card_box.get_sample_ys(30)

    # Determine top and bottom border color
    for edge, start_y, increment_y in [(TOP, 0, +1), (BOTTOM, input_h - 1, -1)]:
        all_color_buckets = {}
        for fixed_x in xs:
            color_buckets = categorize_pixels_in_line(im, (fixed_x, start_y), (0, increment_y), depth_to_calibrate)
            all_color_buckets[fixed_x] = color_buckets
        borders.inform_color(edge, all_color_buckets)

    # Determine left and right border color
    for edge, start_x, increment_x in [(LEFT, 0, +1), (RIGHT, input_w - 1, -1)]:
        all_color_buckets = {}
        for fixed_y in ys:
            color_buckets = categorize_pixels_in_line(im, (start_x, fixed_y), (increment_x, 0), depth_to_calibrate)
            all_color_buckets[fixed_y] = color_buckets
        borders.inform_color(edge, all_color_buckets)

    # Determine top and bottom border thickness
    for edge, start, increment in [(TOP, 0, +1), (BOTTOM, input_h - 1, -1)]:
        all_depths = {}
        color = borders.color[edge]
        if color in (Borders.SOLID_BLACK, Borders.SOLID_WHITE):
            for fixed_x in xs:
                depth = detect_depth_in_column(im, color, fixed_x, start, increment, max_depth_to_check)
                all_depths[fixed_x] = depth
                if SAVE_INTERMEDIATE_FOR_DEBUG:
                    draw.line([(fixed_x, start), (fixed_x, start + increment * depth)], fill='magenta')

            borders.inform_depth(edge, all_depths)

            if SAVE_INTERMEDIATE_FOR_DEBUG:
                depth_calculated = borders.thickness[edge]
                border_y_calculated = start + increment * depth_calculated
                draw.line([(0, border_y_calculated), (input_w - 1, border_y_calculated)], fill='magenta')

    if SAVE_INTERMEDIATE_FOR_DEBUG:
        path = make_intermediate_name(loaded_input_card.path, 'step0bordercheck')
        im.save(path, 'JPEG', dpi=(dpi, dpi), quality=100, subsampling=0)

    return borders