#!/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)