shape_converter.py

from collections import defaultdict
from PIL import Image
import numpy as np
import sys


WHITE = 0


class UnionFind:
    def __init__(self, n):
        self.id = [i for i in range(n)]

    def root(self, i):
        while(i != self.id[i]):
            self.id[i] = self.id[self.id[i]]
            i = self.id[i]
        return i

    def union(self, p, q):
        i = self.root(p)
        j = self.root(q)

        if i != j:
            self.id[j] = i
            return True

        return False


def image_to_matrix(path):
    im = Image.open(path).convert('L')
    w, h = im.size
    binary = im.point(lambda p: p > 128 and 1)
    binary = binary.resize((w//2,h//2),Image.NEAREST)
    w, h = binary.size
    return np.array(binary)

def is_legal_move(matrix, i, j):
    return i >= 0 and j >= 0 and i < len(matrix) and j < len(matrix[i]) and matrix[i][j] == WHITE

def get_id(row, col, n_col):
    return row * n_col + col
    
def detect_union(matrix):
    h, w = len(matrix), len(matrix[0])
    uf = UnionFind(w * h)
    shapes = 0
    visited = set()
    for i in range(h):
        for j in range(w):
            if (matrix[i][j]) == WHITE:
                shapes += 1

    for i in range(h):
        for j in range(w):
            if matrix[i][j] == WHITE and (i,j) not in visited:
                visited.add((i,j))
                for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                    nx = i + dx
                    ny = j + dy
                    if is_legal_move(matrix, nx, ny):
                        id_cur = get_id(i, j, w)
                        id_new = get_id(nx, ny, w) 
                        if uf.union(id_cur, id_new):
                            shapes -= 1

    print(f"shapes {shapes}")
    return uf


if len(sys.argv) < 2:
    print(f"{sys.argv[0]} <image_path>")
    exit(-1)

path = sys.argv[1]
matrix = image_to_matrix(path)

# print matrix
for r in range(len(matrix)):
    for c in range(len(matrix[r])):
        print(matrix[r,c],end='')
    print()


union = detect_union(matrix)

h = len(matrix)
w = len(matrix[0])

print(f"{h}x{w}")

parts = defaultdict(list)
for r in range(h):
    c = 0
    while c < w:
        if matrix[r][c] == WHITE:
            start = (r,c)
            while c < w and matrix[r][c] == WHITE :
                c += 1
            end = (r, c-1)
            id = union.root(get_id(start[0], start[1], w))
            parts[id].append(start)
            parts[id].append(end)

        c += 1

for part in parts:
    print(sorted(parts[part]))
print(len(parts))