diff --git a/map2gazebo/map2gazebo_offline.py b/map2gazebo/map2gazebo_offline.py index 585ab0e..6493232 100644 --- a/map2gazebo/map2gazebo_offline.py +++ b/map2gazebo/map2gazebo_offline.py @@ -7,65 +7,75 @@ import os import sys -class MapConverter(): + +class MapConverter: def __init__(self, map_dir, export_dir, threshold=105, height=2.0): - + self.threshold = threshold self.height = height self.export_dir = export_dir self.map_dir = map_dir def map_callback(self): - + map_array = cv2.imread(self.map_dir) map_array = cv2.flip(map_array, 0) - print(f'loading map file: {self.map_dir}') + print(f"loading map file: {self.map_dir}") try: map_array = cv2.cvtColor(map_array, cv2.COLOR_BGR2GRAY) except cv2.error as err: - print(err, "Conversion failed: Invalid image input, please check your file path") + print( + err, + "Conversion failed: Invalid image input, please check your file path", + ) sys.exit() - info_dir = self.map_dir.replace('pgm','yaml') + info_dir = self.map_dir.replace("pgm", "yaml") + + with open(info_dir, "r") as stream: + map_info = yaml.load(stream, Loader=yaml.FullLoader) - with open(info_dir, 'r') as stream: - map_info = yaml.load(stream, Loader=yaml.FullLoader) - # set all -1 (unknown) values to 0 (unoccupied) map_array[map_array < 0] = 0 contours = self.get_occupied_regions(map_array) - print('Processing...') + print("Processing...") meshes = [self.contour_to_mesh(c, map_info) for c in contours] corners = list(np.vstack(contours)) corners = [c[0] for c in corners] mesh = trimesh.util.concatenate(meshes) - if not self.export_dir.endswith('/'): - self.export_dir = self.export_dir + '/' - file_dir = self.export_dir + map_info['image'].replace('pgm','stl') - print(f'export file: {file_dir}') - - with open(file_dir, 'wb') as f: + filename = os.path.basename(map_info["image"]).replace("pgm", "stl") + file_dir = os.path.join(self.export_dir, filename) + print(f"export file: {file_dir}") + + # Ensure export directory exists + os.makedirs(self.export_dir, exist_ok=True) + + with open(file_dir, "wb") as f: mesh.export(f, "stl") - + def get_occupied_regions(self, map_array): """ Get occupied regions of map """ map_array = map_array.astype(np.uint8) - _, thresh_map = cv2.threshold( - map_array, self.threshold, 100, cv2.THRESH_BINARY) + _, thresh_map = cv2.threshold(map_array, self.threshold, 100, cv2.THRESH_BINARY) contours, hierarchy = cv2.findContours( - thresh_map, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE) + thresh_map, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE + ) # Using cv2.RETR_CCOMP classifies external contours at top level of - # hierarchy and interior contours at second level. + # hierarchy and interior contours at second level. # If the whole space is enclosed by walls RETR_EXTERNAL will exclude # all interior obstacles e.g. furniture. # https://docs.opencv.org/trunk/d9/d8b/tutorial_py_contours_hierarchy.html hierarchy = hierarchy[0] output_contours = [] for idx, contour in enumerate(contours): - output_contours.append(contour) if 0 not in contour else print('Remove image boundary') - + ( + output_contours.append(contour) + if 0 not in contour + else print("Remove image boundary") + ) + return output_contours def contour_to_mesh(self, contour, metadata): @@ -75,57 +85,62 @@ def contour_to_mesh(self, contour, metadata): x, y = point[0] vertices = [] new_vertices = [ - self.coords_to_loc((x, y), metadata), - self.coords_to_loc((x, y+1), metadata), - self.coords_to_loc((x+1, y), metadata), - self.coords_to_loc((x+1, y+1), metadata)] + self.coords_to_loc((x, y), metadata), + self.coords_to_loc((x, y + 1), metadata), + self.coords_to_loc((x + 1, y), metadata), + self.coords_to_loc((x + 1, y + 1), metadata), + ] vertices.extend(new_vertices) vertices.extend([v + height for v in new_vertices]) - faces = [[0, 2, 4], - [4, 2, 6], - [1, 2, 0], - [3, 2, 1], - [5, 0, 4], - [1, 0, 5], - [3, 7, 2], - [7, 6, 2], - [7, 4, 6], - [5, 4, 7], - [1, 5, 3], - [7, 3, 5]] + faces = [ + [0, 2, 4], + [4, 2, 6], + [1, 2, 0], + [3, 2, 1], + [5, 0, 4], + [1, 0, 5], + [3, 7, 2], + [7, 6, 2], + [7, 4, 6], + [5, 4, 7], + [1, 5, 3], + [7, 3, 5], + ] mesh = trimesh.Trimesh(vertices=vertices, faces=faces) if not mesh.is_volume: mesh.fix_normals() meshes.append(mesh) mesh = trimesh.util.concatenate(meshes) - mesh.remove_duplicate_faces() + mesh.update_faces(mesh.unique_faces()) # mesh will still have internal faces. Would be better to get # all duplicate faces and remove both of them, since duplicate faces # are guaranteed to be internal faces return mesh - def coords_to_loc(self,coords, metadata): + def coords_to_loc(self, coords, metadata): x, y = coords - loc_x = x * metadata['resolution'] + metadata['origin'][0] - loc_y = y * metadata['resolution'] + metadata['origin'][1] + loc_x = x * metadata["resolution"] + metadata["origin"][0] + loc_y = y * metadata["resolution"] + metadata["origin"][1] # TODO: transform (x*res, y*res, 0.0) by Pose map_metadata.origin # instead of assuming origin is at z=0 with no rotation wrt map frame return np.array([loc_x, loc_y, 0.0]) + if __name__ == "__main__": parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS) parser.add_argument( - '--map_dir', type=str, required=True, - help='File name of the map to convert' + "--map_dir", type=str, required=True, help="File name of the map to convert" ) parser.add_argument( - '--export_dir', type=str, default=os.path.abspath('.'), - help='Mesh output directory' + "--export_dir", + type=str, + default=os.path.abspath("."), + help="Mesh output directory", ) option = parser.parse_args() Converter = MapConverter(option.map_dir, option.export_dir) Converter.map_callback() - print('Conversion Done') + print("Conversion Done")