Tracing boundaries in raster images

piccells.py

from __future__ import annotations
from typing import NamedTuple, Iterable, Tuple, Sequence, Optional, Set, AbstractSet, List, Dict, Iterator, Mapping
from dataclasses import dataclass
from io import StringIO
from collections import defaultdict
from enum import Enum, auto as enum_auto

#       (x, y)   — → — (x,   y,   XINC) — → — (x+1, y)
#         |                                      |
#         ↑                                      ↓
#         |                                      |
#   (x, y+1, YDEC)         (x, y)          (x+1, y, YINC)
#         |                                      |
#         ↑                                      ↓
#         |                                      |
#       (x, y+1) — ← — (x+1, y+1, XDEC) — ← — (x+1, y+1)

@dataclass(frozen=True)
class Z2:
    __slots__ = ('x', 'y')
    x: int
    y: int

    def __add__(self, other: Z2) -> Z2:
        return Z2(self.x + other.x, self.y + other.y)

    def __sub__(self, other: Z2) -> Z2:
        return Z2(self.x - other.x, self.y - other.y)

    def __mul__(self, other: int) -> Z2:
        return Z2(self.x * other, self.y * other)

    def __rmul__(self, other: int) -> Z2:
        return Z2(self.x * other, self.y * other)

    @property
    def seq(self) -> Sequence[int]:
        return self.x, self.y

    def __repr__(self) -> str:
        return self.seq.__repr__()

Z2_00, Z2_10, Z2_11, Z2_01 = Z2(0, 0), Z2(1, 0), Z2(1, 1), Z2(0, 1)
Z2_N0, Z2_0N, Z2_NN = Z2(-1, 0), Z2(0, -1), Z2(-1, -1)

@dataclass(frozen=True)
class Cell: # A pixel’s interior
    __slots__ = ('c',)
    c: Z2

    def boundary(self, orientation: bool = True) -> List[Edge]:
        raw_edges = [Edge(self.c,         EdgeOr.XINC),
                     Edge(self.c + Z2_10, EdgeOr.YINC),
                     Edge(self.c + Z2_11, EdgeOr.XDEC),
                     Edge(self.c + Z2_01, EdgeOr.YDEC)]
        if orientation:
            return raw_edges
        raw_edges.reverse()
        return [e.reverse for e in raw_edges]

    def __repr__(self) -> str:
        return f'Cell{self.c.seq}'

class EdgeOrData(NamedTuple):
    target_delta: Z2
    outsider_delta: Z2
    index_: int

class EdgeOr(Enum):
    value: EdgeOrData
    XINC = EdgeOrData(Z2_10, Z2_0N, 0)
    YINC = EdgeOrData(Z2_01, Z2_00, 1)
    XDEC = EdgeOrData(Z2_N0, Z2_N0, 2)
    YDEC = EdgeOrData(Z2_0N, Z2_NN, 3)

    @property
    def target_delta(self) -> Z2:
        return self.value.target_delta

    @property
    def outsider_delta(self) -> Z2:
        return self.value.outsider_delta

    @property
    def reverse(self) -> EdgeOr:
        return _EDGE_ORS[self.value.index_ + 2] # pylint: disable=no-member

    @property
    def rotate(self) -> EdgeOr:
        return _EDGE_ORS[self.value.index_ + 1] # pylint: disable=no-member

    def __repr__(self) -> str:
        return self.name

_EDGE_ORS: Sequence[EdgeOr] = tuple(EdgeOr) * 2

@dataclass(frozen=True)
class Edge:
    __slots__ = ('src', 'ori')
    src: Z2
    ori: EdgeOr

    @property
    def reverse(self) -> Edge:
        return Edge(self.src + self.ori.target_delta, self.ori.reverse)

    @property
    def source(self) -> Z2:
        return self.src

    @property
    def target(self) -> Z2:
        return self.src + self.ori.target_delta

    @property
    def insider(self) -> Cell:
        return self.reverse.outsider

    @property
    def outsider(self) -> Cell:
        return Cell(self.src + self.ori.outsider_delta)

    def translate(self, amount: int) -> Edge:
        return Edge(self.src + self.ori.target_delta * amount, self.ori)

    def __repr__(self) -> str:
        return f'{self.src.seq}→{self.target.seq}'

@dataclass(frozen=True)
class Path:
    edges: Tuple[Edge, ...]

    @property
    def end_edges(self) -> Optional[Tuple[Edge, Edge]]:
        edges = self.edges
        return (edges[0], edges[-1]) if edges else None

@dataclass(frozen=True)
class Cycle:
    edges: Tuple[Edge, ...]

    def dump_edges(self, target: Set[Edge]) -> AbstractSet[Edge]:
        target.update(self.edges)
        return target

    def dump_svg(self, s: StringIO) -> None:
        edges = self.edges
        if edges:
            edge = edges[0]
            s.write(f'M {edge.source.x} {edge.source.y} ')
        for edge in edges[1:]:
            s.write(f'L {edge.source.x} {edge.source.y} ')
        s.write('z ')

@dataclass
class CellAttrs:
    black: bool = False
    level: Optional[int] = None

    def __init__(self) -> None:
        self.black = False
        self.level = None

class Picture(Mapping[Cell, CellAttrs]):
    _orig_lines: List[str]
    _size: Z2
    _pixels: Dict[Z2, CellAttrs]

    def __init__(self, picture_str: str) -> None:
        lines = self._orig_lines = picture_str.splitlines()
        line_count = len(lines)
        if not line_count:
            raise ValueError('Should have at least one row.')
        col_count = len(lines[0])
        if any(len(line) != col_count for line in lines[1:]):
            raise ValueError('Should have equal column count in each row.')
        pixels = self._pixels = defaultdict(CellAttrs)
        for y, line in enumerate(lines):
            for x, ch in enumerate(line):
                pixels[Z2(x, y)].black = not ch.isspace()
        self._size = Z2(col_count, line_count)

    def str_with_edges(self, edges: AbstractSet[Edge],
                       print_levels: bool = False) -> str:
        arrs_y = {(False, False): '  ', (True, False): ' ↓',
                  (False, True):  ' ↑', (True, True):  ' ↕'}
        arrs_x = {(False, False): '    ', (True, False): '   →',
                  (False, True):  '   ←', (True, True):  '   ↔'}

        def write_y_arrow(s: StringIO, x: int, y: int) -> None:
            edge = Edge(Z2(x, y), EdgeOr.YINC)
            s.write(arrs_y[(edge in edges, edge.reverse in edges)])

        def write_x_arrow(s: StringIO, x: int, y: int) -> None:
            edge = Edge(Z2(x, y), EdgeOr.XINC)
            s.write(arrs_x[(edge in edges, edge.reverse in edges)])

        col_count, line_count = self._size.seq
        with StringIO() as s:
            for y, line in enumerate(self._orig_lines):
                for x in range(col_count):
                    write_x_arrow(s, x, y)
                s.write('\n')
                for x, ch in enumerate(line):
                    write_y_arrow(s, x, y)
                    if print_levels:
                        s.write(format(self._pixels[Z2(x, y)].level, '2'))
                    else:
                        s.write(' ')
                        s.write(ch)
                write_y_arrow(s, col_count, y)
                s.write('\n')
            for x in range(col_count):
                write_x_arrow(s, x, line_count)
            s.write('\n')
            return s.getvalue()

    @property
    def size(self) -> Z2:
        return self._size

    def __getitem__(self, cell: Cell) -> CellAttrs:
        return self._pixels[cell.c]

    def __iter__(self) -> Iterator[Cell]:
        x_count, y_count = self._size.seq
        for y in range(y_count):
            for x in range(x_count):
                yield Cell(Z2(x, y))

    def __len__(self) -> int:
        return len(self._pixels)

    # Stricter than `getitem(cell) is not None`
    def __contains__(self, cell: object) -> bool:
        if not isinstance(cell, Cell):
            return False
        x_count, y_count = self._size.seq
        return 0 <= cell.c.x < x_count and 0 <= cell.c.y < y_count

    def outer_cells(self) -> Iterator[Cell]:
        def iter_line(start: Z2, stop: Z2, delta: Z2) -> Iterator[Cell]:
            c = start
            while c != stop:
                yield Cell(c)
                c += delta

        x_last, y_last = (self._size - Z2_11).seq
        yield from iter_line(Z2(0, 0), Z2(x_last, 0), Z2_10)
        yield from iter_line(Z2(x_last, 0), Z2(x_last, y_last), Z2_01)
        yield from iter_line(Z2(x_last, y_last), Z2(0, y_last), Z2_N0)
        yield from iter_line(Z2(0, y_last), Z2(0, 0), Z2_0N)

class PathVertex(Enum):
    SOURCE = enum_auto()
    TARGET = enum_auto()
    # TRANSIT = enum_auto()

def find_cycles(pic: Picture) -> Iterable[Cycle]:
    vertex_paths: Dict[Z2, Set[Tuple[PathVertex, Path]]]
    vertex_paths = defaultdict(set)
    paths: Set[Path] = set()
    cycles: Set[Cycle] = set()

    def may_connect(first: Edge, second: Edge) -> bool:
        if first.insider != second.insider:
            return True
        cell = first.translate(1).outsider
        cell_ = second.translate(-1).outsider
        assert cell == cell_, 'This should be a square grid.'
        return not pic[cell].black

    def _remove(path: Optional[Path]) -> None:
        if path is None:
            return
        assert (end_edges := path.end_edges) is not None
        first, last = end_edges
        paths.remove(path)
        vertex_paths[first.source].remove((PathVertex.SOURCE, path))
        vertex_paths[last.target].remove((PathVertex.TARGET, path))

    def _add(path: Path) -> None:
        # May be “closed” with `may_connect(path, path) == False`
        assert (end_edges := path.end_edges) is not None
        first, last = end_edges
        paths.add(path)
        vertex_paths[first.source].add((PathVertex.SOURCE, path))
        vertex_paths[last.target].add((PathVertex.TARGET, path))

    def incorporate(new_edge: Edge) -> None:
        source = new_edge.source
        before_source: Optional[Path] = None
        for kind, other_path in vertex_paths[source]:
            if kind is not PathVertex.TARGET:
                continue
            if not may_connect(other_path.edges[-1], new_edge):
                continue
            before_source = other_path
            break

        target = new_edge.target
        after_target: Optional[Path] = None
        for kind, other_path in vertex_paths[target]:
            if kind is not PathVertex.SOURCE:
                continue
            if not may_connect(new_edge, other_path.edges[0]):
                continue
            after_target = other_path
            break

        new_edges = (new_edge,)
        if before_source is None and after_target is None:
            _add(Path(new_edges))
        elif before_source is after_target:
            assert before_source is not None
            _remove(before_source)
            add_cycle(Cycle(before_source.edges + new_edges))
        else:
            _remove(before_source)
            _remove(after_target)
            before = before_source.edges if before_source else ()
            after = after_target.edges if after_target else ()
            _add(Path(before + new_edges + after))

    def add_cycle(cycle: Cycle) -> None:
        cycles.add(cycle)

    boundary: Set[Cell] = set(pic.outer_cells())
    white_boundary = set(cell for cell in boundary if not pic[cell].black)
    activated_cells: Set[Cell]
    level: int
    if white_boundary:
        activated_cells = white_boundary
        level = 0
    else: # Totally black
        activated_cells = boundary
        level = 1
    activate_later: Set[Cell] = set()

    while activated_cells or activate_later:
        if not activated_cells:
            # The previous level is complete
            level += 1
            activated_cells, activate_later = activate_later, set()
        cell = activated_cells.pop()
        attrs = pic[cell]
        attrs.level = level
        for edge in cell.boundary():
            neighbor = edge.outsider
            neighbor_black = pic[neighbor].black
            if neighbor in pic and pic[neighbor].level is None:
                if attrs.black is neighbor_black:
                    activated_cells.add(neighbor)
                elif neighbor not in activated_cells:
                    activate_later.add(neighbor)
            if not attrs.black:
                continue
            if neighbor_black:
                continue
            # Now, this cell is black and the neighbor is white:
            incorporate(edge)

    return cycles

def main() -> None:
    p = ("          D\n"
         "  AAAAAA   \n"
         " AA     AA \n"
         " AA B   AA \n"
         "  A  CC  A \n"
         "  A     AA \n"
         "  AAAAAAA  \n")

    pic = Picture(p)
    cycles = find_cycles(pic)

    def cell_outside(cell: Cell) -> bool:
        return pic[cell].level in (0, None)

    edges: Set[Edge] = set()
    s = StringIO()
    for cycle in cycles:
        any_edge = cycle.edges[0]
        if (cell_outside(any_edge.insider) or
            cell_outside(any_edge.outsider)):
            # Outermost cycles go to hell in this easy manner
            continue
        cycle.dump_edges(edges)
        cycle.dump_svg(s)
    cycles_svg = s.getvalue()
    s.close()

    print('=' * 60)
    print(pic.str_with_edges(frozenset()))
    print('=' * 60)
    print(pic.str_with_edges(edges, print_levels=True))
    print('=' * 60)
    print(cycles_svg)

if __name__ == "__main__":
    main()