initial commit

scripts/skeleton_extract.py

@@ -0,0 +1,79 @@
+#!/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()