80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
skeleton_extract.py
|
|
Extract Manhattan-style skeleton from binary layout images.
|
|
Supports optional color inversion and basic denoising.
|
|
Outputs skeleton PNG and a JSON summary with connected-component counts.
|
|
"""
|
|
import argparse
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from PIL import Image
|
|
import cv2
|
|
from skimage.morphology import skeletonize
|
|
from skimage import img_as_ubyte
|
|
import json
|
|
|
|
|
|
def load_image(path, invert=False):
|
|
im = Image.open(path).convert('L')
|
|
a = np.array(im)
|
|
# auto threshold by Otsu
|
|
_, th = cv2.threshold(a, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
if invert:
|
|
th = 255 - th
|
|
return th // 255
|
|
|
|
|
|
def denoise(bin_img, kernel=3):
|
|
# simple opening/closing
|
|
k = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel, kernel))
|
|
img = cv2.morphologyEx((bin_img * 255).astype('uint8'), cv2.MORPH_OPEN, k)
|
|
img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, k)
|
|
return img // 255
|
|
|
|
|
|
def extract_skeleton(bin_img):
|
|
# skimage expects bool image
|
|
sk = skeletonize(bin_img > 0)
|
|
return sk.astype('uint8')
|
|
|
|
|
|
def save_png(arr, path):
|
|
im = Image.fromarray((arr * 255).astype('uint8'))
|
|
im.save(path)
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('input')
|
|
p.add_argument('outdir')
|
|
p.add_argument('--invert', action='store_true', help='Invert black/white before processing')
|
|
p.add_argument('--denoise', type=int, default=3, help='Denoise kernel size')
|
|
args = p.parse_args()
|
|
|
|
inp = Path(args.input)
|
|
out = Path(args.outdir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
|
|
bin_img = load_image(inp, invert=args.invert)
|
|
if args.denoise and args.denoise > 0:
|
|
bin_img = denoise(bin_img, kernel=args.denoise)
|
|
|
|
sk = extract_skeleton(bin_img)
|
|
|
|
save_png(bin_img, out / (inp.stem + '_bin.png'))
|
|
save_png(sk, out / (inp.stem + '_sk.png'))
|
|
|
|
# summary
|
|
num_pixels = int(sk.sum())
|
|
components = cv2.connectedComponents((sk * 255).astype('uint8'))[0] - 1
|
|
info = {'input': str(inp), 'pixels_in_skeleton': int(num_pixels), 'components': int(components)}
|
|
with open(out / (inp.stem + '_sk.json'), 'w') as f:
|
|
json.dump(info, f, indent=2)
|
|
|
|
print('Saved:', out / (inp.stem + '_sk.png'))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|