parallelize rendering
This commit is contained in:
parent
316cbe065a
commit
2337591f6f
2 changed files with 263 additions and 110 deletions
|
|
@ -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
|
||||
|
|
|
|||
371
src/__init__.py
371
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",
|
||||
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}")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue