158 lines
4.3 KiB
Python
158 lines
4.3 KiB
Python
#!/usr/bin/python
|
|
# -*- coding: utf-8 -*-
|
|
import numpy as np
|
|
import cv2
|
|
|
|
sub_paths = []
|
|
visited = set()
|
|
current_path = []
|
|
sk = None
|
|
|
|
DO_PRINT_DEBUG = False
|
|
|
|
|
|
def clear():
|
|
sub_paths.clear()
|
|
visited.clear()
|
|
current_path.clear()
|
|
|
|
|
|
def get_neighbors(point):
|
|
"""Get the neighboring points of a given point in the skeleton."""
|
|
x, y = point
|
|
neighbors = []
|
|
for dx in [-1, 0, 1]:
|
|
for dy in [-1, 0, 1]:
|
|
if dx == 0 and dy == 0:
|
|
continue
|
|
nx, ny = x + dx, y + dy
|
|
if 0 <= nx < sk.shape[0] and 0 <= ny < sk.shape[1]:
|
|
if sk[nx, ny] == 255: # Check if it's part of the skeleton
|
|
neighbors.append((nx, ny))
|
|
return neighbors
|
|
|
|
|
|
def add_a_point(p):
|
|
if p in visited:
|
|
return
|
|
if DO_PRINT_DEBUG:
|
|
print(f'Adding: {p}')
|
|
visited.add(p)
|
|
current_path.append(p)
|
|
neighbors = [pix for pix in get_neighbors(p) if pix not in visited]
|
|
branch_count = len(neighbors)
|
|
# print(f'LN: {len(neighbors)} {neighbors}')
|
|
while branch_count == 1:
|
|
new_point = neighbors[0]
|
|
visited.add(new_point)
|
|
current_path.append(new_point)
|
|
neighbors = [pix for pix in get_neighbors(new_point) if pix not in visited]
|
|
branch_count = len(neighbors)
|
|
# print(f'LN: {len(neighbors)} {neighbors}')
|
|
|
|
sub_paths.append(current_path.copy())
|
|
current_path.clear()
|
|
for nei in neighbors:
|
|
add_a_point(nei)
|
|
|
|
|
|
def compute_all_the_sub_paths(skeleton):
|
|
global sk
|
|
sk = skeleton
|
|
clear()
|
|
skeleton_points = np.argwhere(skeleton == 255) # Get all skeleton points
|
|
for skp in skeleton_points:
|
|
pto = tuple(skp)
|
|
add_a_point(pto)
|
|
return sub_paths
|
|
|
|
|
|
def draw_paths(image_shape, sub_paths):
|
|
"""Draw the remaining paths on a blank image."""
|
|
drawn_image = np.zeros(image_shape, dtype=np.uint8) # Create a blank image
|
|
for path in sub_paths:
|
|
for point in path:
|
|
drawn_image[point] = 255 # Draw the path in white
|
|
return drawn_image
|
|
|
|
|
|
def give_me_the_main_loop_only(input_skeleton):
|
|
# Extract sub-paths
|
|
sub_paths = compute_all_the_sub_paths(input_skeleton)
|
|
end_points = {}
|
|
|
|
def is_close_enough(p1, p2):
|
|
x, y = p1
|
|
px, py = p2
|
|
dx = abs(x - px)
|
|
dy = abs(y - py)
|
|
if dx <= 3 and dy <= 3:
|
|
return True
|
|
return False
|
|
|
|
def find_close_point(pto):
|
|
for p in end_points.keys():
|
|
if is_close_enough(pto, p):
|
|
return p
|
|
return None
|
|
|
|
def add_endpoint(pto, idx):
|
|
where = find_close_point(pto)
|
|
if where is None:
|
|
# no close point found, add it
|
|
end_points[pto] = []
|
|
where = pto
|
|
end_points[where].append(idx)
|
|
|
|
# Print the results
|
|
to_be_removed = set()
|
|
for i, path in enumerate(sub_paths):
|
|
# if len(path) < 2:
|
|
# continue
|
|
begin = path[0]
|
|
end = path[-1]
|
|
if is_close_enough(begin, end):
|
|
to_be_removed.add(i)
|
|
continue
|
|
add_endpoint(begin, i)
|
|
add_endpoint(end, i)
|
|
if DO_PRINT_DEBUG:
|
|
print(f"Sub-path {i}: {len(path):4} {path[0]},{path[-1]}")
|
|
|
|
# Draw the remaining paths on a blank image
|
|
# to_be_removed = {0, 7, 12, 13}
|
|
# to_be_removed = set()
|
|
for n, (cc, ep) in enumerate(end_points.items()):
|
|
pto = int(cc[0]), int(cc[1])
|
|
if DO_PRINT_DEBUG:
|
|
print(f'Endpoint #{n} {pto}: {len(ep)}, {ep}')
|
|
if len(ep) == 1:
|
|
# this is a dead branch
|
|
to_be_removed.add(ep[0])
|
|
if DO_PRINT_DEBUG:
|
|
print(f'{to_be_removed=}')
|
|
sub_paths_remaining = [sp for n, sp in enumerate(sub_paths) if n not in to_be_removed]
|
|
|
|
return draw_paths(input_skeleton.shape, sub_paths_remaining)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
skeleton_image_path = 'sk_bonk.png' # Update with your image path
|
|
skeleton = cv2.imread(skeleton_image_path, cv2.IMREAD_GRAYSCALE)
|
|
|
|
# Threshold the image to ensure it's binary
|
|
_, skeleton = cv2.threshold(skeleton, 127, 255, cv2.THRESH_BINARY)
|
|
|
|
cv2.imshow('Input', skeleton)
|
|
|
|
# sp = compute_all_the_sub_paths(skeleton=skeleton)
|
|
# print(len(sp))
|
|
# for n, p in enumerate(sp):
|
|
# print(f'{n:3} {len(p)}')
|
|
|
|
loop_only = give_me_the_main_loop_only(skeleton)
|
|
cv2.imshow('Output', loop_only)
|
|
|
|
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|