213 lines
7.4 KiB
Python
213 lines
7.4 KiB
Python
|
#!/usr/bin/python
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
import cv2
|
||
|
|
from pathlib import Path
|
||
|
|
import numpy as np
|
||
|
|
from anytree import Node, RenderTree, PreOrderIter
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from typing import Any
|
||
|
|
from skimage.morphology import skeletonize
|
||
|
|
from visit_points import give_me_the_main_loop_only
|
||
|
|
from yaml_stuff import create_corrected_yaml
|
||
|
|
|
||
|
|
DO_SHOW = False
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class MyContour:
|
||
|
|
idx: int
|
||
|
|
c: cv2.Mat | np.ndarray[Any, np.dtype] | np.ndarray
|
||
|
|
h: list # 0_next, 1_previous, 2_child, 3_parent
|
||
|
|
area: float
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
return f'IDX: {self.idx}, Area:{self.area:.0f} h:{self.h}'
|
||
|
|
|
||
|
|
|
||
|
|
border_size = 4
|
||
|
|
|
||
|
|
|
||
|
|
def collect_my_children(me: Node, cl: list[MyContour], threshold=0):
|
||
|
|
"""
|
||
|
|
adds all the linked children to a father Node.
|
||
|
|
:param me: me, the starting Node
|
||
|
|
:param cl: the complete element list (the CV2 hierarchy)
|
||
|
|
:return:
|
||
|
|
"""
|
||
|
|
# print(f'Collecting: {me}')
|
||
|
|
my_first_child = me.name.h[2]
|
||
|
|
if my_first_child < 0:
|
||
|
|
return # You are having no child => we have done
|
||
|
|
if cl[my_first_child].area >= threshold:
|
||
|
|
# print(f'Adding a Node (first child): {my_first_child}, area={cl[my_first_child].area}')
|
||
|
|
child = Node(cl[my_first_child], parent=me)
|
||
|
|
if child.name.h[2] >= 0:
|
||
|
|
# This child has other children, collect them first
|
||
|
|
collect_my_children(child, cl, threshold=threshold)
|
||
|
|
# Now do iterate over all the siblings
|
||
|
|
sibling_idx = cl[my_first_child].h[0]
|
||
|
|
while sibling_idx >= 0:
|
||
|
|
# there is a sibling
|
||
|
|
# print(f'Checking Sibling: {sibling_idx}')
|
||
|
|
if cl[sibling_idx].area >= threshold:
|
||
|
|
# print(f'Adding a Node (sibling child): {sibling_idx}')
|
||
|
|
sibling = Node(cl[sibling_idx], parent=me)
|
||
|
|
if sibling.name.h[2] >= 0:
|
||
|
|
collect_my_children(sibling, cl, threshold=threshold)
|
||
|
|
sibling_idx = cl[sibling_idx].h[0]
|
||
|
|
|
||
|
|
|
||
|
|
FILTERED = '_filtered'
|
||
|
|
RACE = '_race'
|
||
|
|
COMPOSITE = '_composite'
|
||
|
|
|
||
|
|
BLACK_LIST = (
|
||
|
|
FILTERED,
|
||
|
|
RACE,
|
||
|
|
COMPOSITE,
|
||
|
|
'_edited',
|
||
|
|
'_raceline',
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def process_a_map(in_image):
|
||
|
|
"""
|
||
|
|
Partendo da una file di mappa grezzo, tenta di creare:
|
||
|
|
una versione filtrata per localizzazione (_filtered.pgm)
|
||
|
|
una versione filtrata per calcolo della race-line (_race.pgm)
|
||
|
|
una versione filtrata per visualizzazione e monitoraggio algoritmo (_composite.png)
|
||
|
|
:param in_image:
|
||
|
|
:return:
|
||
|
|
"""
|
||
|
|
if in_image.suffix.lower() != '.pgm':
|
||
|
|
# Processa solo le immagini di tipo pgm
|
||
|
|
return
|
||
|
|
stem = str(in_image.stem)
|
||
|
|
for bl in BLACK_LIST:
|
||
|
|
if stem.endswith(bl):
|
||
|
|
# scarta tutte le immagini che sono degli output di elaborazioni precedenti
|
||
|
|
return
|
||
|
|
# Load the image
|
||
|
|
image_ = cv2.imread(str(in_image), cv2.IMREAD_GRAYSCALE)
|
||
|
|
# print(image_.shape)
|
||
|
|
image = cv2.copyMakeBorder(image_, top=border_size, bottom=border_size, left=border_size,
|
||
|
|
right=border_size, borderType=cv2.BORDER_CONSTANT, value=205)
|
||
|
|
if DO_SHOW:
|
||
|
|
cv2.imshow(in_image.stem, image)
|
||
|
|
color_image = cv2.cvtColor(image_, cv2.COLOR_GRAY2BGR)
|
||
|
|
|
||
|
|
# Apply binary threshold (you can tweak the threshold value if needed)
|
||
|
|
_, thresh = cv2.threshold(image, 210, 255, cv2.THRESH_BINARY)
|
||
|
|
|
||
|
|
# Find contours
|
||
|
|
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
# print('++ LEN ++++++++++++++++++', len(contours))
|
||
|
|
if len(contours) < 1:
|
||
|
|
raise ValueError(f'No contours found for :{in_image}')
|
||
|
|
for c in contours:
|
||
|
|
area = cv2.contourArea(c)
|
||
|
|
# print(f'{area=}')
|
||
|
|
if area < 100: # TODO Hardcodein is not a good practice
|
||
|
|
continue
|
||
|
|
mask_inside = np.zeros_like(image)
|
||
|
|
|
||
|
|
# Draw the contour in white
|
||
|
|
cv2.drawContours(mask_inside, [c], -1, color=255, thickness=cv2.FILLED)
|
||
|
|
# mask_outside = cv2.bitwise_not(mask_inside)
|
||
|
|
|
||
|
|
img_retain = cv2.bitwise_and(image, mask_inside)
|
||
|
|
|
||
|
|
thres = 230
|
||
|
|
_, th2 = cv2.threshold(img_retain, thres, 255, cv2.THRESH_BINARY)
|
||
|
|
|
||
|
|
ccc, hhh = cv2.findContours(th2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
cl: list[MyContour] = []
|
||
|
|
for n, (c, h) in enumerate(zip(ccc, *hhh)):
|
||
|
|
cl.append(MyContour(idx=n, c=c, h=h, area=cv2.contourArea(c)))
|
||
|
|
# print(n, h, cv2.contourArea(c))
|
||
|
|
# do find the root contour:
|
||
|
|
roots = tuple(filter(lambda x: x.h[3] < 0, cl))
|
||
|
|
assert len(roots) == 1, 'We found more than one root contour!!'
|
||
|
|
root = Node(roots[0])
|
||
|
|
main_th = root.name.area * 0.01
|
||
|
|
collect_my_children(root, cl, threshold=main_th)
|
||
|
|
tentative = np.zeros_like(image)
|
||
|
|
cv2.fillPoly(tentative, [root.name.c], color=255) # Fill with white color (255)
|
||
|
|
for node in PreOrderIter(root):
|
||
|
|
# print(node.name.idx)
|
||
|
|
if node.name.idx != root.name.idx:
|
||
|
|
cv2.fillPoly(tentative, [node.name.c], color=0) # Fill with white color (255)
|
||
|
|
|
||
|
|
rev = cv2.bitwise_not(tentative)
|
||
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) # shape and size
|
||
|
|
enlarged_black = cv2.dilate(rev, kernel, iterations=1)
|
||
|
|
smoothed_black = cv2.erode(enlarged_black, kernel, iterations=1)
|
||
|
|
filtered = cv2.bitwise_not(smoothed_black)
|
||
|
|
fc = filtered[4:-4, 4:-4] # [rows, cols] => [y1:y2, x1:x2]
|
||
|
|
if DO_SHOW:
|
||
|
|
cv2.imshow(f'Filtered', fc)
|
||
|
|
|
||
|
|
destination = in_image.parent / f'{in_image.stem}{FILTERED}.pgm'
|
||
|
|
create_corrected_yaml(in_image, destination)
|
||
|
|
cv2.imwrite(str(destination), fc)
|
||
|
|
|
||
|
|
binary_image = fc // 255
|
||
|
|
# Skeletonize the binary image
|
||
|
|
skeleton = skeletonize(binary_image)
|
||
|
|
|
||
|
|
# Convert the skeleton back to a format suitable for display
|
||
|
|
skeleton_image = (skeleton * 255).astype(np.uint8)
|
||
|
|
if DO_SHOW:
|
||
|
|
cv2.imshow(f'SK', skeleton_image)
|
||
|
|
|
||
|
|
main_loop = give_me_the_main_loop_only(skeleton_image)
|
||
|
|
|
||
|
|
race_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (30, 30)) # shape and size
|
||
|
|
race_map_1 = cv2.dilate(main_loop, race_kernel, iterations=1)
|
||
|
|
race_map = cv2.bitwise_and(race_map_1, fc)
|
||
|
|
if DO_SHOW:
|
||
|
|
cv2.imshow(f'Race', race_map)
|
||
|
|
|
||
|
|
destination = in_image.parent / f'{in_image.stem}{RACE}.pgm'
|
||
|
|
create_corrected_yaml(in_image, destination)
|
||
|
|
cv2.imwrite(str(destination), race_map)
|
||
|
|
|
||
|
|
# Define colors for the masks (BGR format)
|
||
|
|
color1 = (0, 255, 0) # Green
|
||
|
|
color2 = (255, 0, 0) # Blue
|
||
|
|
|
||
|
|
# Create colored masks
|
||
|
|
colored_mask1 = np.zeros_like(color_image)
|
||
|
|
colored_mask1[fc == 255] = color1
|
||
|
|
|
||
|
|
colored_mask2 = np.zeros_like(color_image)
|
||
|
|
colored_mask2[race_map == 255] = color2
|
||
|
|
|
||
|
|
# Blend the images
|
||
|
|
sum_image = cv2.bitwise_xor(colored_mask1, colored_mask2)
|
||
|
|
composite_image = cv2.addWeighted(color_image, 0.7, sum_image, 0.2, 0)
|
||
|
|
|
||
|
|
# Show the result
|
||
|
|
# cv2.imshow('Composite Image', composite_image)
|
||
|
|
|
||
|
|
destination = in_image.parent / f'{in_image.stem}{COMPOSITE}.png'
|
||
|
|
create_corrected_yaml(in_image, destination)
|
||
|
|
cv2.imwrite(str(destination), composite_image)
|
||
|
|
|
||
|
|
# cv2.imwrite('comp.png', composite_image)
|
||
|
|
|
||
|
|
if DO_SHOW:
|
||
|
|
cv2.waitKey(0)
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
FOLDER = Path(r'C:\Mc\Python\PyProjs\atlanta_maps\atlant_maps\maps')
|
||
|
|
for element in FOLDER.iterdir():
|
||
|
|
if element.is_dir():
|
||
|
|
img_filename = element / f'{element.stem}.pgm'
|
||
|
|
if img_filename.is_file():
|
||
|
|
print(f'We have found: {img_filename}')
|
||
|
|
process_a_map(img_filename)
|
||
|
|
|