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