parallelize rendering

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

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="<min lon>,<min lat>,<max lon>,<max lat>",
         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",
+        "--square",
-        help="File path for STL output",
+        action="store_true",
-        required=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",
+        "--output",
-        action="store_true",
+        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")
+    map_bounds = tuple(float(coord) for coord in args.bbox.split(","))
-    natural = osm.get_natural()
-    buildings = osm.get_buildings()
-
-    map_bounds = [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 = get_elev_grid(elev_map, map_bounds, args.slices * 25, scale_factor)
-    elev_grid = np.zeros((ELEV_RES, ELEV_RES))
+    elev_grid *= args.elev_exaggeration
-    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
 
-    ROAD_HEIGHT = 0.25
+    renderers = []
-    # Render roads on slopes by over-extruding them and then calculating their
+    for i in range(args.slices):
-    # intersection with a slightly offset elevation grid
+        for j in range(args.slices):
-    drive_mask = None
+            renderers.append(Renderer(
-    for i in range(ELEV_RES):
+                map_bounds=(
-        for j in range(ELEV_RES):
+                    map_bounds[0] + map_width / args.slices * i,
-            if elev_grid[i, j] > 0:
+                    map_bounds[1] + map_height / args.slices * j,
-                model += solid2.cube([
+                    map_bounds[0] + map_width / args.slices * (i + 1),
-                    model_bounds[2] / ELEV_RES,
+                    map_bounds[1] + map_height / args.slices * (j + 1),
-                    model_bounds[3] / ELEV_RES,
+                ),
-                    elev_grid[i, j] + args.baseplate_h,
+                chunk_bounds=(
-                ]).translate(
-                    model_bounds[2] / ELEV_RES * i,
-                    model_bounds[3] / ELEV_RES * j,
                     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,
-                )
+                    model_bounds[2] / args.slices,
-                if drive_mask is None:
+                    model_bounds[3] / args.slices,
-                    drive_mask = drive_mask_cube
+                ),
-                else:
+                scale_factor=scale_factor,
-                    drive_mask += drive_mask_cube
+                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...")
+    print("Combining...")
-    drive_polygons = []
+    chunks = []
-    for shape in drive_net["geometry"]:
+    for i in range(args.slices):
-        shape = convert_shape_from_world_coords(shape).simplify(1).buffer(0.25)
+        for j in range(args.slices):
-        drive_polygons.append(shape.intersection(model_rect))
+            mesh = stl.Mesh.from_file(f"/tmp/chunk.{i}.{j}.stl")
-    drive_polygons = [shape for shape in drive_polygons if not shape.is_empty]
+            mesh.x += model_bounds[2] / args.slices * i
-    for shape in union_all(drive_polygons).geoms:
+            mesh.y += model_bounds[3] / args.slices * j
-        model += solid2.linear_extrude(100)(to_scad_polygon(shape)).translate(0, 0, args.baseplate_h) * drive_mask
+            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...")
+    duration = datetime.now() - t_start
-    for i, building in buildings.iterrows():
+    print(f"Finished in: {duration}")
-        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)