From 2337591f6f4175f8279a8bac52f1c2839c3eeee7 Mon Sep 17 00:00:00 2001 From: Brent Schroeter Date: Sat, 17 May 2025 12:31:24 -0700 Subject: [PATCH] parallelize rendering --- requirements.txt | 2 + src/__init__.py | 371 +++++++++++++++++++++++++++++++++-------------- 2 files changed, 263 insertions(+), 110 deletions(-) diff --git a/requirements.txt b/requirements.txt index 6d706bf..b4d4c96 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ geopandas==1.0.1 gdal[numpy]==3.6.2 +numpy==2.2.5 +numpy-stl==3.2.0 pyrosm==0.6.2 shapely==2.1.0 solidpython2==2.1.1 diff --git a/src/__init__.py b/src/__init__.py index 73c7e02..665fb2c 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,13 +1,20 @@ import argparse +import subprocess +import sys -from math import asin, cos, pi, sin, sqrt +from dataclasses import dataclass +from datetime import datetime +from math import asin, ceil, cos, pi, sin, sqrt +from multiprocessing import cpu_count, Pool +from typing import Tuple import numpy as np import solid2 +import stl from osgeo import gdal from pyrosm import OSM -from shapely import affinity, Polygon, union_all +from shapely import affinity, MultiPolygon, Polygon, union_all def haversine_dist(lat1, lng1, lat2, lng2): @@ -25,6 +32,47 @@ def haversine_dist(lat1, lng1, lat2, lng2): return 2 * EARTH_RADIUS * asin(sqrt(hav_theta)) +def get_elev_grid(elev_map, map_bounds, resolution, scale_factor): + """ + Generate a 2D array of elevation values (in model units), smoothed and + "blocked" to avoid jagged and/or superfluous edges. + """ + elev_grid = np.zeros((resolution, resolution)) + for i in range(resolution): + for j in range(resolution): + elev_grid[i, j] = elev_map.get_elevation( + map_bounds[1] + (map_bounds[3] - map_bounds[1]) / resolution * j, + map_bounds[0] + (map_bounds[2] - map_bounds[0]) / resolution * i, + ) + elev_grid = (elev_grid - elev_grid.min()) * scale_factor + + # Convolve a uniform kernel over matrix to smooth values + KERNEL_SIZE = 5 + kernel_offset = (KERNEL_SIZE - 1) / 2 + assert kernel_offset % 1 == 0 + kernel_offset = int(kernel_offset) + smoothed = np.zeros((resolution, resolution)) + for i in range(resolution): + for j in range(resolution): + # Clamping to the array bounds makes the edges of the model less + # smooth, but that's acceptable for these purposes + windowed = elev_grid[ + max(i - kernel_offset, 0):min(i + kernel_offset, elev_grid.shape[0]), + max(j - kernel_offset, 0):min(j + kernel_offset, elev_grid.shape[1]), + ] + smoothed[i, j] = windowed.mean() + + BLOCK_TOLERANCE = 0.1 + for block_size in range(2, resolution + 1): + for i in range(resolution - block_size + 1): + for j in range(resolution - block_size + 1): + block = smoothed[i:i+block_size, j:j+block_size].flatten() + if block.max() - block.min() < BLOCK_TOLERANCE: + smoothed[i:i+block_size, j:j+block_size] = block.mean() + + return smoothed + + def to_scad_polygon(shapely_polygon): """ Generate a solid2 SCAD polygon object from a shapely Polygon. @@ -41,7 +89,7 @@ def get_building_height(building): """ Infer the height, in meters, of a building from OSM building data. """ - DEFAULT_METERS_PER_LEVEL = 5 + DEFAULT_METERS_PER_LEVEL = 4 DEFAULT_LEVELS_PER_BUILDING = 3 if building["height"] is not None: return float(building["height"]) @@ -69,6 +117,138 @@ class ElevationMap: raise IndexError("coordinates are outside of elevation map area") +def render_stl(model, file_path): + """ + Like solid2.render_to_stl_file(), but with improved argument interpolation + and stdout/stderr piping. + """ + solid2.scad_render_to_file(model, f"{file_path}.scad") + subprocess.check_call( + ["openscad", "-o", file_path, f"{file_path}.scad"], + stdout=sys.stdout, + stderr=sys.stderr, + ) + + +@dataclass +class Renderer: + map_bounds: Tuple[float, float, float, float] + chunk_bounds: Tuple[float, float, float, float] + scale_factor: float + osm_path: str + elev_grid: np.array + baseplate_h: float + output: str + + def render_chunk(self): + chunk_bounds = self.chunk_bounds + map_bounds = self.map_bounds + scale_factor = self.scale_factor + elev_grid = self.elev_grid + baseplate_h = self.baseplate_h + + map_width = map_bounds[2] - map_bounds[0] + map_height = map_bounds[3] - map_bounds[1] + + elev_res = elev_grid.shape[0] + assert elev_res == elev_grid.shape[1] + + osm = OSM(self.osm_path) + solid2.set_global_fn(48) + + drive_net = osm.get_network(network_type="driving+service") + buildings = osm.get_buildings() + + chunk_rect = Polygon([ + (0, 0), + (0, chunk_bounds[3]), + (chunk_bounds[2], chunk_bounds[3]), + (chunk_bounds[2], 0), + ]) + + def convert_shape_from_world_coords(shape): + # Move the coordinate system so that the origin is at the bottom left + # of the map rectangle, then scale the shape in x and y directions + return affinity.scale( + affinity.translate(shape, -map_bounds[0], -map_bounds[1]), + chunk_bounds[2] / map_width, + chunk_bounds[3] / map_height, + origin=(0, 0), + ) + + model = solid2.cube([chunk_bounds[2] - chunk_bounds[0], chunk_bounds[3] - chunk_bounds[1], baseplate_h]) + + ROAD_HEIGHT = 0.25 + # Render roads on slopes by over-extruding them and then calculating their + # intersection with a slightly offset elevation grid + drive_mask = None + for i in range(elev_res): + for j in range(elev_res): + if elev_grid[i, j] > 0: + model += solid2.cube([ + chunk_bounds[2] / elev_res, + chunk_bounds[3] / elev_res, + elev_grid[i, j] + baseplate_h, + ]).translate( + chunk_bounds[2] / elev_res * i, + chunk_bounds[3] / elev_res * j, + 0, + ) + drive_mask_cube = solid2.cube([ + chunk_bounds[2] / elev_res, + chunk_bounds[3] / elev_res, + elev_grid[i, j] + baseplate_h + ROAD_HEIGHT, + ]).translate( + chunk_bounds[2] / elev_res * i, + chunk_bounds[3] / elev_res * j, + 0, + ) + if drive_mask is None: + drive_mask = drive_mask_cube + else: + drive_mask += drive_mask_cube + + print(f"[{self.output}] Constructing networks...") + drive_polygons = [] + for shape in drive_net["geometry"]: + shape = convert_shape_from_world_coords(shape).simplify(1).buffer(0.25) + drive_polygons.append(shape.intersection(chunk_rect)) + drive_polygons = [shape for shape in drive_polygons if not shape.is_empty] + drive_polygons = union_all(drive_polygons) + if drive_polygons.geom_type == "Polygon": + drive_polygons = MultiPolygon([drive_polygons]) + for shape in drive_polygons.geoms: + model += solid2.linear_extrude(100)(to_scad_polygon(shape)).translate(0, 0, baseplate_h) * drive_mask + + print(f"[{self.output}] Constructing buildings...") + for i, building in buildings.iterrows(): + shape = convert_shape_from_world_coords(building["geometry"]) + shape = shape.buffer(0.125) + # Simplifying building shapes after scaling can speed up rendering by + # an order of magnitude + shape = shape.simplify(0.05) + shape = shape.intersection(chunk_rect) + if not shape.is_empty: + height = get_building_height(building) * scale_factor + building_center = building["geometry"].centroid + # Computing elevation from the elevation map can result in weird + # artifacts, so calculate based on the less granular precomputed + # grid instead + elev = elev_grid[ + min(max(int((building_center.x - map_bounds[0]) / map_width * elev_res), 0), elev_res - 1), + min(max(int((building_center.y - map_bounds[1]) / map_height * elev_res), 0), elev_res - 1), + ] + if shape.geom_type == "Polygon": + model += solid2.linear_extrude(args.baseplate_h + elev + ROAD_HEIGHT + height)(to_scad_polygon(shape)) + else: + for shape in shape.geoms: + model += solid2.linear_extrude(args.baseplate_h + elev + ROAD_HEIGHT + height)(to_scad_polygon(shape)) + + print(f"[{self.output}] Rendering...") + render_stl(model, self.output) + print(f"[{self.output}] Chunk finished.") + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -86,6 +266,18 @@ if __name__ == "__main__": help=",,,", required=True, ) + parser.add_argument( + "--elev-exaggeration", + type=float, + help="Additional scaling factor to apply to elevations for dramatic effect", + default=1.0, + ) + parser.add_argument( + "--slices", + type=int, + help="Parallelize render and increase elevation resolution by splitting model into n tiles per side (--slices=3 splits work into 9 tiles)", + default=ceil(sqrt(cpu_count())), + ) parser.add_argument( "--edge-len", type=float, @@ -93,9 +285,9 @@ if __name__ == "__main__": help="Long edge length of output STL bounding box", ) parser.add_argument( - "--output", - help="File path for STL output", - required=True, + "--square", + action="store_true", + help="Constrains the model to a square with dimension --edge-len", ) parser.add_argument( "--baseplate-h", @@ -104,20 +296,17 @@ if __name__ == "__main__": help="Depth of the base plate in millimeters", ) parser.add_argument( - "--circular", - action="store_true", + "--output", + help="File path for STL output", + required=True, ) args = parser.parse_args() + t_start = datetime.now() + elev_map = ElevationMap(args.topo) - osm = OSM(args.osm) - solid2.set_global_fn(48) - drive_net = osm.get_network(network_type="driving") - natural = osm.get_natural() - buildings = osm.get_buildings() - - map_bounds = [float(coord) for coord in args.bbox.split(",")] + map_bounds = tuple(float(coord) for coord in args.bbox.split(",")) map_width = map_bounds[2] - map_bounds[0] map_height = map_bounds[3] - map_bounds[1] @@ -139,6 +328,22 @@ if __name__ == "__main__": map_width_meters = map_width * meters_per_unit_lng map_height_meters = map_height * meters_per_unit_lat + + if args.square: + map_center = ( + (map_bounds[0] + map_bounds[2]) / 2, + (map_bounds[1] + map_bounds[3]) / 2, + ) + short_edge_meters = min(map_width_meters, map_height_meters) + map_bounds = ( + map_center[0] - short_edge_meters / meters_per_unit_lng / 2, + map_center[1] - short_edge_meters / meters_per_unit_lat / 2, + map_center[0] + short_edge_meters / meters_per_unit_lng / 2, + map_center[1] + short_edge_meters / meters_per_unit_lat / 2, + ) + map_width_meters = short_edge_meters + map_height_meters = short_edge_meters + aspect_ratio = map_width_meters / map_height_meters if aspect_ratio > 1: # Wide @@ -146,12 +351,6 @@ if __name__ == "__main__": else: # Tall model_bounds = (0, 0, args.edge_len * aspect_ratio, args.edge_len) - model_rect = Polygon([ - (0, 0), - (0, model_bounds[3]), - (model_bounds[2], model_bounds[3]), - (model_bounds[2], 0), - ]) # Scale factor from real-world meters to the model coordinate system scale_factor = model_bounds[2] / map_width_meters @@ -166,99 +365,51 @@ if __name__ == "__main__": f"{scale_factor:.4g}", ) - def convert_shape_from_world_coords(shape): - # Move the coordinate system so that the origin is at the bottom left - # of the map rectangle, then scale the shape in x and y directions - return affinity.scale( - affinity.translate(shape, -map_bounds[0], -map_bounds[1]), - model_bounds[2] / map_width, - model_bounds[3] / map_height, - origin=(0, 0), - ) - - model = solid2.cube([model_bounds[2] - model_bounds[0], model_bounds[3] - model_bounds[1], args.baseplate_h]) - print("Calculating elevations...") # Slice the model area into a rectilinear grid to approximate topography - ELEV_RES = 30 - elev_grid = np.zeros((ELEV_RES, ELEV_RES)) - for i in range(ELEV_RES): - for j in range(ELEV_RES): - elev_grid[i, j] = elev_map.get_elevation( - map_bounds[1] + (map_bounds[3] - map_bounds[1]) / ELEV_RES * j, - map_bounds[0] + (map_bounds[2] - map_bounds[0]) / ELEV_RES * i, - ) - min_world_elev = elev_grid.min() - elev_grid = (elev_grid - min_world_elev) * scale_factor + elev_grid = get_elev_grid(elev_map, map_bounds, args.slices * 25, scale_factor) + elev_grid *= args.elev_exaggeration - ROAD_HEIGHT = 0.25 - # Render roads on slopes by over-extruding them and then calculating their - # intersection with a slightly offset elevation grid - drive_mask = None - for i in range(ELEV_RES): - for j in range(ELEV_RES): - if elev_grid[i, j] > 0: - model += solid2.cube([ - model_bounds[2] / ELEV_RES, - model_bounds[3] / ELEV_RES, - elev_grid[i, j] + args.baseplate_h, - ]).translate( - model_bounds[2] / ELEV_RES * i, - model_bounds[3] / ELEV_RES * j, + renderers = [] + for i in range(args.slices): + for j in range(args.slices): + renderers.append(Renderer( + map_bounds=( + map_bounds[0] + map_width / args.slices * i, + map_bounds[1] + map_height / args.slices * j, + map_bounds[0] + map_width / args.slices * (i + 1), + map_bounds[1] + map_height / args.slices * (j + 1), + ), + chunk_bounds=( 0, - ) - drive_mask_cube = solid2.cube([ - model_bounds[2] / ELEV_RES, - model_bounds[3] / ELEV_RES, - elev_grid[i, j] + args.baseplate_h + ROAD_HEIGHT, - ]).translate( - model_bounds[2] / ELEV_RES * i, - model_bounds[3] / ELEV_RES * j, 0, - ) - if drive_mask is None: - drive_mask = drive_mask_cube - else: - drive_mask += drive_mask_cube + model_bounds[2] / args.slices, + model_bounds[3] / args.slices, + ), + scale_factor=scale_factor, + osm_path=args.osm, + elev_grid=elev_grid[ + int(elev_grid.shape[0]/args.slices*i):int(elev_grid.shape[0]/args.slices*(i+1)), + int(elev_grid.shape[1]/args.slices*j):int(elev_grid.shape[1]/args.slices*(j+1)), + ], + baseplate_h=args.baseplate_h, + output=f"/tmp/chunk.{i}.{j}.stl", + )) + with Pool(cpu_count()) as pool: + pool.map(Renderer.render_chunk, renderers) - print("Constructing networks...") - drive_polygons = [] - for shape in drive_net["geometry"]: - shape = convert_shape_from_world_coords(shape).simplify(1).buffer(0.25) - drive_polygons.append(shape.intersection(model_rect)) - drive_polygons = [shape for shape in drive_polygons if not shape.is_empty] - for shape in union_all(drive_polygons).geoms: - model += solid2.linear_extrude(100)(to_scad_polygon(shape)).translate(0, 0, args.baseplate_h) * drive_mask + print("Combining...") + chunks = [] + for i in range(args.slices): + for j in range(args.slices): + mesh = stl.Mesh.from_file(f"/tmp/chunk.{i}.{j}.stl") + mesh.x += model_bounds[2] / args.slices * i + mesh.y += model_bounds[3] / args.slices * j + chunks.append(mesh) + # This doesn't actually perform a union calculation but rather lines up + # multiple meshes alongside each other, which as long as it's good enough + # for a gcode slicer... + stl.Mesh(np.concatenate([chunk.data for chunk in chunks])).save(args.output, mode=stl.Mode.BINARY) - print("Constructing buildings...") - for i, building in buildings.iterrows(): - shape = convert_shape_from_world_coords(building["geometry"]) - shape = shape.buffer(0.125) - # Simplifying building shapes after scaling can speed up rendering by - # an order of magnitude - shape = shape.simplify(0.05) - shape = shape.intersection(model_rect) - if not shape.is_empty: - height = get_building_height(building) * scale_factor - building_center = building["geometry"].centroid - # Computing elevation from the elevation map can result in weird - # artifacts, so calculate based on the less granular precomputed - # grid instead - elev = elev_grid[ - min(max(int((building_center.x - map_bounds[0]) / map_width * ELEV_RES), 0), ELEV_RES - 1), - min(max(int((building_center.y - map_bounds[1]) / map_height * ELEV_RES), 0), ELEV_RES - 1), - ] - if shape.geom_type == "Polygon": - model += solid2.linear_extrude(args.baseplate_h + elev + ROAD_HEIGHT + height)(to_scad_polygon(shape)) - else: - for shape in shape.geoms: - model += solid2.linear_extrude(args.baseplate_h + elev + ROAD_HEIGHT + height)(to_scad_polygon(shape)) - - if args.circular: - radius = min(model_bounds[2], model_bounds[3]) / 2 - model *= solid2.translate(model_bounds[2] / 2, model_bounds[3] / 2, 0)( - solid2.cylinder(100, r=radius), - ) - - print("Rendering... (grab a cup of coffee)") - solid2.render_to_stl_file(model, args.output) + duration = datetime.now() - t_start + print(f"Finished in: {duration}")