voxcity 0.6.2__py3-none-any.whl → 0.6.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of voxcity might be problematic. Click here for more details.

@@ -1,791 +1,791 @@
1
- import numpy as np
2
- import os
3
- import trimesh
4
- import matplotlib.colors as mcolors
5
- import matplotlib.cm as cm
6
- import matplotlib.pyplot as plt
7
-
8
- def create_voxel_mesh(voxel_array, class_id, meshsize=1.0, building_id_grid=None, mesh_type=None):
9
- """
10
- Create a 3D mesh from voxels preserving sharp edges, scaled by meshsize.
11
-
12
- This function converts a 3D voxel array into a triangulated mesh, where each voxel
13
- face is converted into two triangles. The function preserves sharp edges between
14
- different classes and handles special cases for buildings.
15
-
16
- Parameters
17
- ----------
18
- voxel_array : np.ndarray (3D)
19
- The voxel array of shape (X, Y, Z) where each cell contains a class ID.
20
- - 0: typically represents void/air
21
- - -2: typically represents trees
22
- - -3: typically represents buildings
23
- Other values can represent different classes as defined by the application.
24
-
25
- class_id : int
26
- The ID of the class to extract. Only voxels with this ID will be included
27
- in the output mesh.
28
-
29
- meshsize : float, default=1.0
30
- The real-world size of each voxel in meters, applied uniformly to x, y, and z
31
- dimensions. Used to scale the output mesh to real-world coordinates.
32
-
33
- building_id_grid : np.ndarray (2D), optional
34
- 2D grid of building IDs with shape (X, Y). Only used when class_id=-3 (buildings).
35
- Each cell contains a unique identifier for the building at that location.
36
- This allows tracking which faces belong to which building.
37
-
38
- mesh_type : str, optional
39
- Type of mesh to create, controlling which faces are included:
40
- - None (default): create faces at boundaries between different classes
41
- - 'building_solar': only create faces at boundaries between buildings (-3)
42
- and either void (0) or trees (-2). Useful for solar analysis
43
- where only exposed surfaces matter.
44
-
45
- Returns
46
- -------
47
- mesh : trimesh.Trimesh or None
48
- The resulting triangulated mesh for the given class_id. Returns None if no
49
- voxels of the specified class are found.
50
-
51
- The mesh includes:
52
- - vertices: 3D coordinates of each vertex
53
- - faces: triangles defined by vertex indices
54
- - face_normals: normal vectors for each face
55
- - metadata: If class_id=-3, includes 'building_id' mapping faces to buildings
56
-
57
- Examples
58
- --------
59
- Basic usage for a simple voxel array:
60
- >>> voxels = np.zeros((10, 10, 10))
61
- >>> voxels[4:7, 4:7, 0:5] = 1 # Create a simple column
62
- >>> mesh = create_voxel_mesh(voxels, class_id=1, meshsize=0.5)
63
-
64
- Creating a building mesh with IDs:
65
- >>> building_ids = np.zeros((10, 10))
66
- >>> building_ids[4:7, 4:7] = 1 # Mark building #1
67
- >>> mesh = create_voxel_mesh(voxels, class_id=-3,
68
- ... building_id_grid=building_ids,
69
- ... meshsize=1.0)
70
-
71
- Notes
72
- -----
73
- - The function creates faces only at boundaries between different classes or at
74
- the edges of the voxel array.
75
- - Each face is split into two triangles for compatibility with graphics engines.
76
- - Face normals are computed to ensure correct lighting and rendering.
77
- - For buildings (class_id=-3), building IDs are tracked to maintain building identity.
78
- - The mesh preserves sharp edges, which is important for architectural visualization.
79
- """
80
- # Find voxels of the current class
81
- voxel_coords = np.argwhere(voxel_array == class_id)
82
-
83
- if building_id_grid is not None:
84
- building_id_grid_flipud = np.flipud(building_id_grid)
85
-
86
- if len(voxel_coords) == 0:
87
- return None
88
-
89
- # Define the 6 faces of a unit cube (local coordinates 0..1)
90
- unit_faces = np.array([
91
- # Front
92
- [[0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]],
93
- # Back
94
- [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]],
95
- # Right
96
- [[1, 0, 0], [1, 1, 0], [1, 1, 1], [1, 0, 1]],
97
- # Left
98
- [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]],
99
- # Top
100
- [[0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 1, 0]],
101
- # Bottom
102
- [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]]
103
- ])
104
-
105
- # Define face normals
106
- face_normals = np.array([
107
- [0, 0, 1], # Front
108
- [0, 0, -1], # Back
109
- [1, 0, 0], # Right
110
- [-1, 0, 0], # Left
111
- [0, 1, 0], # Top
112
- [0, -1, 0] # Bottom
113
- ])
114
-
115
- vertices = []
116
- faces = []
117
- face_normals_list = []
118
- building_ids = [] # List to store building IDs for each face
119
-
120
- for x, y, z in voxel_coords:
121
- # For buildings, get the building ID from the grid
122
- building_id = None
123
- if class_id == -3 and building_id_grid is not None:
124
- building_id = building_id_grid_flipud[x, y]
125
-
126
- # Check each face of the current voxel
127
- adjacent_coords = [
128
- (x, y, z+1), # Front
129
- (x, y, z-1), # Back
130
- (x+1, y, z), # Right
131
- (x-1, y, z), # Left
132
- (x, y+1, z), # Top
133
- (x, y-1, z) # Bottom
134
- ]
135
-
136
- # Only create faces where there's a transition based on mesh_type
137
- for face_idx, adj_coord in enumerate(adjacent_coords):
138
- try:
139
- # If adj_coord is outside array bounds, it's a boundary => face is visible
140
- if adj_coord[0] < 0 or adj_coord[1] < 0 or adj_coord[2] < 0:
141
- is_boundary = True
142
- else:
143
- adj_value = voxel_array[adj_coord]
144
-
145
- if mesh_type == 'open_air' and class_id == -3:
146
- # For building_solar, only create faces at boundaries with void (0) or trees (-2)
147
- is_boundary = (adj_value == 0 or adj_value == -2)
148
- else:
149
- # Default behavior - create faces at any class change
150
- is_boundary = (adj_value == 0 or adj_value != class_id)
151
- except IndexError:
152
- # Out of range => boundary
153
- is_boundary = True
154
-
155
- if is_boundary:
156
- # Local face in (0..1) for x,y,z, then shift by voxel coords
157
- face_verts = (unit_faces[face_idx] + np.array([x, y, z])) * meshsize
158
- current_vert_count = len(vertices)
159
-
160
- vertices.extend(face_verts)
161
- # Convert quad to two triangles
162
- faces.extend([
163
- [current_vert_count, current_vert_count + 1, current_vert_count + 2],
164
- [current_vert_count, current_vert_count + 2, current_vert_count + 3]
165
- ])
166
- # Add face normals for both triangles
167
- face_normals_list.extend([face_normals[face_idx], face_normals[face_idx]])
168
-
169
- # Store building ID for both triangles if this is a building
170
- if class_id == -3 and building_id_grid is not None:
171
- building_ids.extend([building_id, building_id])
172
-
173
- if not vertices:
174
- return None
175
-
176
- vertices = np.array(vertices)
177
- faces = np.array(faces)
178
- face_normals_list = np.array(face_normals_list)
179
-
180
- # Create mesh
181
- mesh = trimesh.Trimesh(
182
- vertices=vertices,
183
- faces=faces,
184
- face_normals=face_normals_list
185
- )
186
-
187
- # Merge vertices that are at the same position
188
- mesh.merge_vertices()
189
-
190
- # Ensure metadata dict exists
191
- if not hasattr(mesh, 'metadata') or mesh.metadata is None:
192
- mesh.metadata = {}
193
-
194
- # Store intended per-triangle normals to avoid reliance on auto-computed normals
195
- mesh.metadata['provided_face_normals'] = face_normals_list
196
-
197
- # Add building IDs as metadata for buildings
198
- if class_id == -3 and building_id_grid is not None and building_ids:
199
- mesh.metadata['building_id'] = np.array(building_ids)
200
-
201
- return mesh
202
-
203
- def create_sim_surface_mesh(sim_grid, dem_grid,
204
- meshsize=1.0, z_offset=1.5,
205
- cmap_name='viridis',
206
- vmin=None, vmax=None):
207
- """
208
- Create a colored planar surface mesh from simulation data, positioned above a DEM.
209
-
210
- This function generates a 3D visualization mesh for 2D simulation results (like
211
- Green View Index, solar radiation, etc.). The mesh is positioned above the Digital
212
- Elevation Model (DEM) by a specified offset, and colored according to the simulation
213
- values using a matplotlib colormap.
214
-
215
- Parameters
216
- ----------
217
- sim_grid : 2D np.ndarray
218
- 2D array of simulation values (e.g., Green View Index, solar radiation).
219
- NaN values in this grid will be skipped in the output mesh.
220
- The grid should be oriented with north at the top.
221
-
222
- dem_grid : 2D np.ndarray
223
- 2D array of ground elevations in meters. Must have the same shape as sim_grid.
224
- Used to position the visualization mesh at the correct height above terrain.
225
-
226
- meshsize : float, default=1.0
227
- Size of each cell in meters. Applied uniformly to x and y dimensions.
228
- Determines the resolution of the output mesh.
229
-
230
- z_offset : float, default=1.5
231
- Additional height offset in meters added to dem_grid elevations.
232
- Used to position the visualization above ground level for better visibility.
233
-
234
- cmap_name : str, default='viridis'
235
- Matplotlib colormap name used for coloring the mesh based on sim_grid values.
236
- Common options:
237
- - 'viridis': Default, perceptually uniform, colorblind-friendly
238
- - 'RdYlBu': Red-Yellow-Blue, good for diverging data
239
- - 'jet': Rainbow colormap (not recommended for scientific visualization)
240
-
241
- vmin : float, optional
242
- Minimum value for color mapping. If None, uses min of sim_grid (excluding NaN).
243
- Used to control the range of the colormap.
244
-
245
- vmax : float, optional
246
- Maximum value for color mapping. If None, uses max of sim_grid (excluding NaN).
247
- Used to control the range of the colormap.
248
-
249
- Returns
250
- -------
251
- mesh : trimesh.Trimesh or None
252
- A single mesh containing one colored square face (two triangles) per non-NaN cell.
253
- Returns None if there are no valid (non-NaN) cells in sim_grid.
254
-
255
- The mesh includes:
256
- - vertices: 3D coordinates of each vertex
257
- - faces: triangles defined by vertex indices
258
- - face_colors: RGBA colors for each face based on sim_grid values
259
- - visual: trimesh.visual.ColorVisuals object storing the face colors
260
-
261
- Examples
262
- --------
263
- Basic usage with Green View Index data:
264
- >>> gvi = np.array([[0.5, 0.6], [0.4, 0.8]]) # GVI values
265
- >>> dem = np.array([[10.0, 10.2], [9.8, 10.1]]) # Ground heights
266
- >>> mesh = create_sim_surface_mesh(gvi, dem, meshsize=1.0, z_offset=1.5)
267
-
268
- Custom color range and colormap:
269
- >>> mesh = create_sim_surface_mesh(gvi, dem,
270
- ... cmap_name='RdYlBu',
271
- ... vmin=0.0, vmax=1.0)
272
-
273
- Notes
274
- -----
275
- - The function automatically creates a matplotlib colorbar figure for visualization
276
- - Both input grids are flipped vertically to match the voxel_array orientation
277
- - Each grid cell is converted to two triangles for compatibility with 3D engines
278
- - The mesh is positioned at dem_grid + z_offset to float above the terrain
279
- - Face colors are interpolated from the colormap based on sim_grid values
280
- """
281
- # Flip arrays vertically
282
- sim_grid_flipped = np.flipud(sim_grid)
283
- dem_grid_flipped = np.flipud(dem_grid)
284
-
285
- # Identify valid (non-NaN) values
286
- valid_mask = ~np.isnan(sim_grid_flipped)
287
- valid_values = sim_grid_flipped[valid_mask]
288
- if valid_values.size == 0:
289
- return None
290
-
291
- # If vmin/vmax not provided, use actual min/max of the valid sim data
292
- if vmin is None:
293
- vmin = np.nanmin(valid_values)
294
- if vmax is None:
295
- vmax = np.nanmax(valid_values)
296
-
297
- # Prepare the colormap and create colorbar
298
- norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
299
- scalar_map = cm.ScalarMappable(norm=norm, cmap=cmap_name)
300
-
301
- # Create a figure just for the colorbar
302
- fig, ax = plt.subplots(figsize=(6, 1))
303
- plt.colorbar(scalar_map, cax=ax, orientation='horizontal')
304
- plt.tight_layout()
305
- plt.close()
306
-
307
- vertices = []
308
- faces = []
309
- face_colors = []
310
-
311
- vert_index = 0
312
- nrows, ncols = sim_grid_flipped.shape
313
-
314
- # Build a quad (two triangles) for each valid cell
315
- for x in range(nrows):
316
- for y in range(ncols):
317
- val = sim_grid_flipped[x, y]
318
- if np.isnan(val):
319
- continue
320
-
321
- z_base = meshsize * int(dem_grid_flipped[x, y] / meshsize + 1.5) + z_offset
322
-
323
- # 4 corners in (x,y)*meshsize
324
- v0 = [ x * meshsize, y * meshsize, z_base ]
325
- v1 = [(x + 1) * meshsize, y * meshsize, z_base ]
326
- v2 = [(x + 1) * meshsize, (y + 1) * meshsize, z_base ]
327
- v3 = [ x * meshsize, (y + 1) * meshsize, z_base ]
328
-
329
- vertices.extend([v0, v1, v2, v3])
330
- faces.extend([
331
- [vert_index, vert_index + 1, vert_index + 2],
332
- [vert_index, vert_index + 2, vert_index + 3]
333
- ])
334
-
335
- # Get color from colormap
336
- color_rgba = np.array(scalar_map.to_rgba(val)) # shape (4,)
337
-
338
- # Each cell has 2 faces => add the color twice
339
- face_colors.append(color_rgba)
340
- face_colors.append(color_rgba)
341
-
342
- vert_index += 4
343
-
344
- if len(vertices) == 0:
345
- return None
346
-
347
- vertices = np.array(vertices, dtype=float)
348
- faces = np.array(faces, dtype=int)
349
- face_colors = np.array(face_colors, dtype=float)
350
-
351
- mesh = trimesh.Trimesh(
352
- vertices=vertices,
353
- faces=faces,
354
- face_colors=face_colors,
355
- process=False # skip auto merge if you want to preserve quads
356
- )
357
-
358
- return mesh
359
-
360
- def create_city_meshes(
361
- voxel_array,
362
- vox_dict,
363
- meshsize=1.0,
364
- include_classes=None,
365
- exclude_classes=None,
366
- ):
367
- """
368
- Create a collection of colored 3D meshes representing different city elements.
369
-
370
- This function processes a voxelized city model and creates separate meshes for
371
- different urban elements (buildings, trees, etc.), each with its own color.
372
- The function preserves sharp edges and applies appropriate colors based on the
373
- provided color dictionary.
374
-
375
- Parameters
376
- ----------
377
- voxel_array : np.ndarray (3D)
378
- 3D array representing the voxelized city model. Each voxel contains a class ID
379
- that maps to an urban element type:
380
- - 0: Void/air (automatically skipped)
381
- - -2: Trees
382
- - -3: Buildings
383
- Other values can represent different urban elements as defined in vox_dict.
384
-
385
- vox_dict : dict
386
- Dictionary mapping class IDs to RGB colors. Each entry should be:
387
- {class_id: [R, G, B]} where R, G, B are 0-255 integer values.
388
- Example: {-3: [200, 200, 200], -2: [0, 255, 0]} for grey buildings and
389
- green trees. The key 0 (air) is automatically excluded.
390
-
391
- meshsize : float, default=1.0
392
- Size of each voxel in meters, applied uniformly to x, y, and z dimensions.
393
- Used to scale the output meshes to real-world coordinates.
394
-
395
- Returns
396
- -------
397
- meshes : dict
398
- Dictionary mapping class IDs to their corresponding trimesh.Trimesh objects.
399
- Each mesh includes:
400
- - vertices: 3D coordinates scaled by meshsize
401
- - faces: triangulated faces preserving sharp edges
402
- - face_colors: RGBA colors from vox_dict
403
- - visual: trimesh.visual.ColorVisuals object storing the face colors
404
-
405
- Classes with no voxels are automatically excluded from the output.
406
-
407
- Examples
408
- --------
409
- Basic usage with buildings and trees:
410
- >>> voxels = np.zeros((10, 10, 10))
411
- >>> voxels[4:7, 4:7, 0:5] = -3 # Add a building
412
- >>> voxels[2:4, 2:4, 0:3] = -2 # Add some trees
413
- >>> colors = {
414
- ... -3: [200, 200, 200], # Grey buildings
415
- ... -2: [0, 255, 0] # Green trees
416
- ... }
417
- >>> meshes = create_city_meshes(voxels, colors, meshsize=1.0)
418
-
419
- Notes
420
- -----
421
- - The function automatically skips class_id=0 (typically air/void)
422
- - Each urban element type gets its own separate mesh for efficient rendering
423
- - Colors are converted from RGB [0-255] to RGBA [0-1] format
424
- - Sharp edges are preserved to maintain architectural features
425
- - Empty classes (no voxels) are automatically excluded from the output
426
- - Errors during mesh creation for a class are caught and reported
427
- """
428
- meshes = {}
429
-
430
- # Convert RGB colors to hex for material properties
431
- color_dict = {k: mcolors.rgb2hex([v[0]/255, v[1]/255, v[2]/255])
432
- for k, v in vox_dict.items() if k != 0} # Exclude air
433
-
434
- # Determine which classes to process
435
- unique_classes = np.unique(voxel_array)
436
-
437
- if include_classes is not None:
438
- # Only keep classes explicitly requested (and present in the data)
439
- class_iterable = [c for c in include_classes if c in unique_classes]
440
- else:
441
- class_iterable = list(unique_classes)
442
-
443
- exclude_set = set(exclude_classes) if exclude_classes is not None else set()
444
-
445
- # Create vertices and faces for each object class
446
- for class_id in class_iterable:
447
- if class_id == 0: # Skip air
448
- continue
449
-
450
- if class_id in exclude_set:
451
- # Explicitly skipped (e.g., will be replaced with custom mesh)
452
- continue
453
-
454
- try:
455
- mesh = create_voxel_mesh(voxel_array, class_id, meshsize=meshsize)
456
-
457
- if mesh is None:
458
- continue
459
-
460
- # Convert hex color to RGBA
461
- if class_id not in color_dict:
462
- # Color not provided; skip silently for robustness
463
- continue
464
- rgb_color = np.array(mcolors.hex2color(color_dict[class_id]))
465
- rgba_color = np.concatenate([rgb_color, [1.0]])
466
-
467
- # Assign color to all faces
468
- mesh.visual.face_colors = np.tile(rgba_color, (len(mesh.faces), 1))
469
-
470
- meshes[class_id] = mesh
471
-
472
- except ValueError as e:
473
- print(f"Skipping class {class_id}: {e}")
474
-
475
- return meshes
476
-
477
- def export_meshes(meshes, output_directory, base_filename):
478
- """
479
- Export a collection of meshes to both OBJ (with MTL) and STL formats.
480
-
481
- This function exports meshes in two ways:
482
- 1. A single combined OBJ file with materials (and associated MTL file)
483
- 2. Separate STL files for each mesh, named with their class IDs
484
-
485
- Parameters
486
- ----------
487
- meshes : dict
488
- Dictionary mapping class IDs to trimesh.Trimesh objects.
489
- Each mesh should have:
490
- - vertices: 3D coordinates
491
- - faces: triangulated faces
492
- - face_colors: RGBA colors (if using materials)
493
-
494
- output_directory : str
495
- Directory path where the output files will be saved.
496
- Will be created if it doesn't exist.
497
-
498
- base_filename : str
499
- Base name for the output files (without extension).
500
- Will be used to create:
501
- - {base_filename}.obj : Combined mesh with materials
502
- - {base_filename}.mtl : Material definitions for OBJ
503
- - {base_filename}_{class_id}.stl : Individual STL files
504
-
505
- Returns
506
- -------
507
- None
508
- Files are written directly to the specified output directory.
509
-
510
- Examples
511
- --------
512
- >>> meshes = {
513
- ... -3: building_mesh, # Building mesh with grey color
514
- ... -2: tree_mesh # Tree mesh with green color
515
- ... }
516
- >>> export_meshes(meshes, 'output/models', 'city_model')
517
-
518
- This will create:
519
- - output/models/city_model.obj
520
- - output/models/city_model.mtl
521
- - output/models/city_model_-3.stl
522
- - output/models/city_model_-2.stl
523
-
524
- Notes
525
- -----
526
- - OBJ/MTL format preserves colors and materials but is more complex
527
- - STL format is simpler but doesn't support colors
528
- - STL files are exported separately for each class for easier processing
529
- - The OBJ file combines all meshes while preserving their materials
530
- - File extensions are automatically added to the base filename
531
- """
532
- # Export combined mesh as OBJ with materials
533
- combined_mesh = trimesh.util.concatenate(list(meshes.values()))
534
- combined_mesh.export(f"{output_directory}/{base_filename}.obj")
535
-
536
- # Export individual meshes as STL
537
- for class_id, mesh in meshes.items():
538
- # Convert class_id to a string for filename
539
- mesh.export(f"{output_directory}/{base_filename}_{class_id}.stl")
540
-
541
- def split_vertices_manual(mesh):
542
- """
543
- Split a mesh into independent faces by duplicating shared vertices.
544
-
545
- This function imitates trimesh's split_vertices() functionality but ensures
546
- complete face independence by giving each face its own copy of vertices.
547
- This is particularly useful for rendering applications where smooth shading
548
- between faces is undesirable, such as architectural visualization in Rhino.
549
-
550
- Parameters
551
- ----------
552
- mesh : trimesh.Trimesh
553
- Input mesh to split. Should have:
554
- - vertices: array of vertex coordinates
555
- - faces: array of vertex indices forming triangles
556
- - visual: Optional ColorVisuals object with face colors
557
-
558
- Returns
559
- -------
560
- out_mesh : trimesh.Trimesh
561
- New mesh where each face is completely independent, with:
562
- - Duplicated vertices for each face
563
- - No vertex sharing between faces
564
- - Preserved face colors if present in input
565
- - Each face as a separate component
566
-
567
- Examples
568
- --------
569
- Basic usage:
570
- >>> vertices = np.array([[0,0,0], [1,0,0], [1,1,0], [0,1,0]])
571
- >>> faces = np.array([[0,1,2], [0,2,3]]) # Two triangles sharing vertices
572
- >>> mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
573
- >>> split_mesh = split_vertices_manual(mesh)
574
- >>> print(f"Original vertices: {len(mesh.vertices)}") # 4 vertices
575
- >>> print(f"Split vertices: {len(split_mesh.vertices)}") # 6 vertices
576
-
577
- With face colors:
578
- >>> colors = np.array([[255,0,0,255], [0,255,0,255]]) # Red and green faces
579
- >>> mesh.visual = trimesh.visual.ColorVisuals(mesh, face_colors=colors)
580
- >>> split_mesh = split_vertices_manual(mesh) # Colors are preserved
581
-
582
- Notes
583
- -----
584
- - Each output face has exactly 3 unique vertices
585
- - Face colors are preserved in the output mesh
586
- - Useful for:
587
- - Preventing smooth shading artifacts
588
- - Ensuring face color independence
589
- - Preparing meshes for CAD software
590
- - Creating sharp edges in architectural models
591
- - Memory usage increases as vertices are duplicated
592
- """
593
- new_meshes = []
594
-
595
- # For each face, build a small, one-face mesh
596
- for face_idx, face in enumerate(mesh.faces):
597
- face_coords = mesh.vertices[face]
598
-
599
- # Create mini-mesh without colors first
600
- mini_mesh = trimesh.Trimesh(
601
- vertices=face_coords,
602
- faces=[[0, 1, 2]],
603
- process=False # skip merging/cleaning
604
- )
605
-
606
- # If the mesh has per-face colors, set the face color properly
607
- if (mesh.visual.face_colors is not None
608
- and len(mesh.visual.face_colors) == len(mesh.faces)):
609
- # Create a visual object with the face color (for one face)
610
- face_color = mesh.visual.face_colors[face_idx]
611
- color_visual = trimesh.visual.ColorVisuals(
612
- mesh=mini_mesh,
613
- face_colors=np.array([face_color]), # One face, one color
614
- vertex_colors=None
615
- )
616
- mini_mesh.visual = color_visual
617
-
618
- new_meshes.append(mini_mesh)
619
-
620
- # Concatenate all the single-face meshes
621
- out_mesh = trimesh.util.concatenate(new_meshes)
622
- return out_mesh
623
-
624
- def save_obj_from_colored_mesh(meshes, output_path, base_filename, max_materials=None):
625
- """
626
- Memory-safe OBJ/MTL exporter.
627
- - Streams vertices/faces to disk (no concatenate, no per-face mini-meshes).
628
- - Uses face colors -> materials (no vertex splitting).
629
- - Optional color quantization to reduce material count.
630
- """
631
- import os
632
- import numpy as np
633
-
634
- os.makedirs(output_path, exist_ok=True)
635
- obj_path = os.path.join(output_path, f"{base_filename}.obj")
636
- mtl_path = os.path.join(output_path, f"{base_filename}.mtl")
637
-
638
- # --------------- helpers ---------------
639
- def to_uint8_rgba(arr):
640
- arr = np.asarray(arr)
641
- if arr.dtype != np.uint8:
642
- # Handle float [0..1] or int [0..255]
643
- if arr.dtype.kind == 'f':
644
- arr = np.clip(arr, 0.0, 1.0)
645
- arr = (arr * 255.0 + 0.5).astype(np.uint8)
646
- else:
647
- arr = arr.astype(np.uint8)
648
- if arr.shape[1] == 3:
649
- alpha = np.full((arr.shape[0], 1), 255, dtype=np.uint8)
650
- arr = np.concatenate([arr, alpha], axis=1)
651
- return arr
652
-
653
- # First pass: build material palette
654
- # We avoid collecting all colors at once—scan per mesh and update a dict.
655
- color_to_id = {}
656
- ordered_colors = [] # list of RGBA uint8 tuples in material order
657
-
658
- # Optional quantizer (lazy-init)
659
- quantizer = None
660
- if max_materials is not None:
661
- try:
662
- from sklearn.cluster import MiniBatchKMeans
663
- quantizer = MiniBatchKMeans(n_clusters=max_materials, random_state=42, batch_size=8192)
664
- # Partial-fit streaming pass over colors
665
- for m in meshes.values():
666
- fc = getattr(m.visual, "face_colors", None)
667
- if fc is None:
668
- continue
669
- fc = to_uint8_rgba(fc)
670
- if fc.size == 0:
671
- continue
672
- # Use only RGB for clustering
673
- quantizer.partial_fit(fc[:, :3].astype(np.float32))
674
- except ImportError:
675
- raise ImportError("scikit-learn is required for color quantization. Install it with: pip install scikit-learn")
676
-
677
- # Assign material ids during a second scan, but still streaming to avoid big unions
678
- def get_material_id(rgba):
679
- key = (int(rgba[0]), int(rgba[1]), int(rgba[2]), int(rgba[3]))
680
- mid = color_to_id.get(key)
681
- if mid is None:
682
- mid = len(ordered_colors)
683
- color_to_id[key] = mid
684
- ordered_colors.append(key)
685
- return mid
686
-
687
- # 2nd pass if quantizing: we need color centroids
688
- centers_u8 = None
689
- if quantizer is not None:
690
- centers = quantizer.cluster_centers_.astype(np.float32) # RGB float
691
- centers = np.clip(centers, 0.0, 255.0).astype(np.uint8)
692
- # Build a quick LUT fun
693
- def quantize_rgb(rgb_u8):
694
- # rgb_u8: (N,3) uint8 -> labels -> centers
695
- labels = quantizer.predict(rgb_u8.astype(np.float32))
696
- return centers[labels]
697
- # We'll convert each mesh's face colors to quantized RGB on the fly
698
- centers_u8 = centers
699
-
700
- # Build materials palette by scanning once (still O(total faces) but tiny memory)
701
- for m in meshes.values():
702
- fc = getattr(m.visual, "face_colors", None)
703
- if fc is None:
704
- # No colors: assign default grey
705
- rgba = np.array([[200,200,200,255]], dtype=np.uint8)
706
- get_material_id(rgba[0])
707
- continue
708
- fc = to_uint8_rgba(fc)
709
- if quantizer is not None:
710
- q_rgb = quantize_rgb(fc[:, :3])
711
- fc = np.concatenate([q_rgb, fc[:, 3:4]], axis=1)
712
- # Iterate unique colors in this mesh to limit get_material_id calls
713
- # but don't materialize huge sets; unique per mesh is fine.
714
- uniq = np.unique(fc, axis=0)
715
- for rgba in uniq:
716
- get_material_id(rgba)
717
-
718
- # Write MTL
719
- with open(mtl_path, "w") as mtl:
720
- for i, (r, g, b, a) in enumerate(ordered_colors):
721
- mtl.write(f"newmtl material_{i}\n")
722
- # Match viewport look: diffuse only, no specular. Many viewers assume sRGB.
723
- kd_r, kd_g, kd_b = r/255.0, g/255.0, b/255.0
724
- mtl.write(f"Kd {kd_r:.6f} {kd_g:.6f} {kd_b:.6f}\n")
725
- # Ambient same as diffuse to avoid darkening in some viewers
726
- mtl.write(f"Ka {kd_r:.6f} {kd_g:.6f} {kd_b:.6f}\n")
727
- # No specular highlight
728
- mtl.write("Ks 0.000000 0.000000 0.000000\n")
729
- # Disable lighting model with specular; keep simple shading
730
- mtl.write("illum 1\n")
731
- # Alpha
732
- mtl.write(f"d {a/255.0:.6f}\n\n")
733
-
734
- # Stream OBJ
735
- with open(obj_path, "w") as obj:
736
- obj.write(f"mtllib {os.path.basename(mtl_path)}\n")
737
-
738
- v_offset = 0 # running vertex index offset
739
- # Reusable cache so we don't keep writing 'usemtl' for the same block unnecessarily
740
- current_material = None
741
-
742
- for class_id, m in meshes.items():
743
- verts = np.asarray(m.vertices, dtype=np.float64)
744
- faces = np.asarray(m.faces, dtype=np.int64)
745
- if verts.size == 0 or faces.size == 0:
746
- continue
747
-
748
- # Write vertices
749
- # (We do a single pass; writing text is the bottleneck, but memory-safe.)
750
- for v in verts:
751
- obj.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
752
-
753
- # Prepare face colors (face-level)
754
- fc = getattr(m.visual, "face_colors", None)
755
- if fc is None or len(fc) != len(faces):
756
- # default grey if missing or mismatched
757
- fc = np.tile(np.array([200,200,200,255], dtype=np.uint8), (len(faces), 1))
758
- else:
759
- fc = to_uint8_rgba(fc)
760
-
761
- if quantizer is not None:
762
- q_rgb = quantize_rgb(fc[:, :3])
763
- fc = np.concatenate([q_rgb, fc[:, 3:4]], axis=1)
764
-
765
- # Group faces by material id and stream in order
766
- # Build material id per face quickly
767
- # Convert face colors to material ids
768
- # (Avoid Python loops over faces more than once)
769
- # Map unique colors in this mesh to material ids first:
770
- uniq_colors, inv_idx = np.unique(fc, axis=0, return_inverse=True)
771
- color_to_mid_local = {tuple(c.tolist()): get_material_id(c) for c in uniq_colors}
772
- mids = np.fromiter(
773
- (color_to_mid_local[tuple(c.tolist())] for c in uniq_colors[inv_idx]),
774
- dtype=np.int64,
775
- count=len(inv_idx)
776
- )
777
-
778
- # Write faces grouped by material, but preserve simple ordering
779
- # Cheap approach: emit runs; switching material only when necessary
780
- current_material = None
781
- for i_face, face in enumerate(faces):
782
- mid = int(mids[i_face])
783
- if current_material != mid:
784
- obj.write(f"usemtl material_{mid}\n")
785
- current_material = mid
786
- a, b, c = face + 1 + v_offset # OBJ is 1-based
787
- obj.write(f"f {a} {b} {c}\n")
788
-
789
- v_offset += len(verts)
790
-
1
+ import numpy as np
2
+ import os
3
+ import trimesh
4
+ import matplotlib.colors as mcolors
5
+ import matplotlib.cm as cm
6
+ import matplotlib.pyplot as plt
7
+
8
+ def create_voxel_mesh(voxel_array, class_id, meshsize=1.0, building_id_grid=None, mesh_type=None):
9
+ """
10
+ Create a 3D mesh from voxels preserving sharp edges, scaled by meshsize.
11
+
12
+ This function converts a 3D voxel array into a triangulated mesh, where each voxel
13
+ face is converted into two triangles. The function preserves sharp edges between
14
+ different classes and handles special cases for buildings.
15
+
16
+ Parameters
17
+ ----------
18
+ voxel_array : np.ndarray (3D)
19
+ The voxel array of shape (X, Y, Z) where each cell contains a class ID.
20
+ - 0: typically represents void/air
21
+ - -2: typically represents trees
22
+ - -3: typically represents buildings
23
+ Other values can represent different classes as defined by the application.
24
+
25
+ class_id : int
26
+ The ID of the class to extract. Only voxels with this ID will be included
27
+ in the output mesh.
28
+
29
+ meshsize : float, default=1.0
30
+ The real-world size of each voxel in meters, applied uniformly to x, y, and z
31
+ dimensions. Used to scale the output mesh to real-world coordinates.
32
+
33
+ building_id_grid : np.ndarray (2D), optional
34
+ 2D grid of building IDs with shape (X, Y). Only used when class_id=-3 (buildings).
35
+ Each cell contains a unique identifier for the building at that location.
36
+ This allows tracking which faces belong to which building.
37
+
38
+ mesh_type : str, optional
39
+ Type of mesh to create, controlling which faces are included:
40
+ - None (default): create faces at boundaries between different classes
41
+ - 'building_solar': only create faces at boundaries between buildings (-3)
42
+ and either void (0) or trees (-2). Useful for solar analysis
43
+ where only exposed surfaces matter.
44
+
45
+ Returns
46
+ -------
47
+ mesh : trimesh.Trimesh or None
48
+ The resulting triangulated mesh for the given class_id. Returns None if no
49
+ voxels of the specified class are found.
50
+
51
+ The mesh includes:
52
+ - vertices: 3D coordinates of each vertex
53
+ - faces: triangles defined by vertex indices
54
+ - face_normals: normal vectors for each face
55
+ - metadata: If class_id=-3, includes 'building_id' mapping faces to buildings
56
+
57
+ Examples
58
+ --------
59
+ Basic usage for a simple voxel array:
60
+ >>> voxels = np.zeros((10, 10, 10))
61
+ >>> voxels[4:7, 4:7, 0:5] = 1 # Create a simple column
62
+ >>> mesh = create_voxel_mesh(voxels, class_id=1, meshsize=0.5)
63
+
64
+ Creating a building mesh with IDs:
65
+ >>> building_ids = np.zeros((10, 10))
66
+ >>> building_ids[4:7, 4:7] = 1 # Mark building #1
67
+ >>> mesh = create_voxel_mesh(voxels, class_id=-3,
68
+ ... building_id_grid=building_ids,
69
+ ... meshsize=1.0)
70
+
71
+ Notes
72
+ -----
73
+ - The function creates faces only at boundaries between different classes or at
74
+ the edges of the voxel array.
75
+ - Each face is split into two triangles for compatibility with graphics engines.
76
+ - Face normals are computed to ensure correct lighting and rendering.
77
+ - For buildings (class_id=-3), building IDs are tracked to maintain building identity.
78
+ - The mesh preserves sharp edges, which is important for architectural visualization.
79
+ """
80
+ # Find voxels of the current class
81
+ voxel_coords = np.argwhere(voxel_array == class_id)
82
+
83
+ if building_id_grid is not None:
84
+ building_id_grid_flipud = np.flipud(building_id_grid)
85
+
86
+ if len(voxel_coords) == 0:
87
+ return None
88
+
89
+ # Define the 6 faces of a unit cube (local coordinates 0..1)
90
+ unit_faces = np.array([
91
+ # Front
92
+ [[0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]],
93
+ # Back
94
+ [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]],
95
+ # Right
96
+ [[1, 0, 0], [1, 1, 0], [1, 1, 1], [1, 0, 1]],
97
+ # Left
98
+ [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]],
99
+ # Top
100
+ [[0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 1, 0]],
101
+ # Bottom
102
+ [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]]
103
+ ])
104
+
105
+ # Define face normals
106
+ face_normals = np.array([
107
+ [0, 0, 1], # Front
108
+ [0, 0, -1], # Back
109
+ [1, 0, 0], # Right
110
+ [-1, 0, 0], # Left
111
+ [0, 1, 0], # Top
112
+ [0, -1, 0] # Bottom
113
+ ])
114
+
115
+ vertices = []
116
+ faces = []
117
+ face_normals_list = []
118
+ building_ids = [] # List to store building IDs for each face
119
+
120
+ for x, y, z in voxel_coords:
121
+ # For buildings, get the building ID from the grid
122
+ building_id = None
123
+ if class_id == -3 and building_id_grid is not None:
124
+ building_id = building_id_grid_flipud[x, y]
125
+
126
+ # Check each face of the current voxel
127
+ adjacent_coords = [
128
+ (x, y, z+1), # Front
129
+ (x, y, z-1), # Back
130
+ (x+1, y, z), # Right
131
+ (x-1, y, z), # Left
132
+ (x, y+1, z), # Top
133
+ (x, y-1, z) # Bottom
134
+ ]
135
+
136
+ # Only create faces where there's a transition based on mesh_type
137
+ for face_idx, adj_coord in enumerate(adjacent_coords):
138
+ try:
139
+ # If adj_coord is outside array bounds, it's a boundary => face is visible
140
+ if adj_coord[0] < 0 or adj_coord[1] < 0 or adj_coord[2] < 0:
141
+ is_boundary = True
142
+ else:
143
+ adj_value = voxel_array[adj_coord]
144
+
145
+ if mesh_type == 'open_air' and class_id == -3:
146
+ # For building_solar, only create faces at boundaries with void (0) or trees (-2)
147
+ is_boundary = (adj_value == 0 or adj_value == -2)
148
+ else:
149
+ # Default behavior - create faces at any class change
150
+ is_boundary = (adj_value == 0 or adj_value != class_id)
151
+ except IndexError:
152
+ # Out of range => boundary
153
+ is_boundary = True
154
+
155
+ if is_boundary:
156
+ # Local face in (0..1) for x,y,z, then shift by voxel coords
157
+ face_verts = (unit_faces[face_idx] + np.array([x, y, z])) * meshsize
158
+ current_vert_count = len(vertices)
159
+
160
+ vertices.extend(face_verts)
161
+ # Convert quad to two triangles
162
+ faces.extend([
163
+ [current_vert_count, current_vert_count + 1, current_vert_count + 2],
164
+ [current_vert_count, current_vert_count + 2, current_vert_count + 3]
165
+ ])
166
+ # Add face normals for both triangles
167
+ face_normals_list.extend([face_normals[face_idx], face_normals[face_idx]])
168
+
169
+ # Store building ID for both triangles if this is a building
170
+ if class_id == -3 and building_id_grid is not None:
171
+ building_ids.extend([building_id, building_id])
172
+
173
+ if not vertices:
174
+ return None
175
+
176
+ vertices = np.array(vertices)
177
+ faces = np.array(faces)
178
+ face_normals_list = np.array(face_normals_list)
179
+
180
+ # Create mesh
181
+ mesh = trimesh.Trimesh(
182
+ vertices=vertices,
183
+ faces=faces,
184
+ face_normals=face_normals_list
185
+ )
186
+
187
+ # Merge vertices that are at the same position
188
+ mesh.merge_vertices()
189
+
190
+ # Ensure metadata dict exists
191
+ if not hasattr(mesh, 'metadata') or mesh.metadata is None:
192
+ mesh.metadata = {}
193
+
194
+ # Store intended per-triangle normals to avoid reliance on auto-computed normals
195
+ mesh.metadata['provided_face_normals'] = face_normals_list
196
+
197
+ # Add building IDs as metadata for buildings
198
+ if class_id == -3 and building_id_grid is not None and building_ids:
199
+ mesh.metadata['building_id'] = np.array(building_ids)
200
+
201
+ return mesh
202
+
203
+ def create_sim_surface_mesh(sim_grid, dem_grid,
204
+ meshsize=1.0, z_offset=1.5,
205
+ cmap_name='viridis',
206
+ vmin=None, vmax=None):
207
+ """
208
+ Create a colored planar surface mesh from simulation data, positioned above a DEM.
209
+
210
+ This function generates a 3D visualization mesh for 2D simulation results (like
211
+ Green View Index, solar radiation, etc.). The mesh is positioned above the Digital
212
+ Elevation Model (DEM) by a specified offset, and colored according to the simulation
213
+ values using a matplotlib colormap.
214
+
215
+ Parameters
216
+ ----------
217
+ sim_grid : 2D np.ndarray
218
+ 2D array of simulation values (e.g., Green View Index, solar radiation).
219
+ NaN values in this grid will be skipped in the output mesh.
220
+ The grid should be oriented with north at the top.
221
+
222
+ dem_grid : 2D np.ndarray
223
+ 2D array of ground elevations in meters. Must have the same shape as sim_grid.
224
+ Used to position the visualization mesh at the correct height above terrain.
225
+
226
+ meshsize : float, default=1.0
227
+ Size of each cell in meters. Applied uniformly to x and y dimensions.
228
+ Determines the resolution of the output mesh.
229
+
230
+ z_offset : float, default=1.5
231
+ Additional height offset in meters added to dem_grid elevations.
232
+ Used to position the visualization above ground level for better visibility.
233
+
234
+ cmap_name : str, default='viridis'
235
+ Matplotlib colormap name used for coloring the mesh based on sim_grid values.
236
+ Common options:
237
+ - 'viridis': Default, perceptually uniform, colorblind-friendly
238
+ - 'RdYlBu': Red-Yellow-Blue, good for diverging data
239
+ - 'jet': Rainbow colormap (not recommended for scientific visualization)
240
+
241
+ vmin : float, optional
242
+ Minimum value for color mapping. If None, uses min of sim_grid (excluding NaN).
243
+ Used to control the range of the colormap.
244
+
245
+ vmax : float, optional
246
+ Maximum value for color mapping. If None, uses max of sim_grid (excluding NaN).
247
+ Used to control the range of the colormap.
248
+
249
+ Returns
250
+ -------
251
+ mesh : trimesh.Trimesh or None
252
+ A single mesh containing one colored square face (two triangles) per non-NaN cell.
253
+ Returns None if there are no valid (non-NaN) cells in sim_grid.
254
+
255
+ The mesh includes:
256
+ - vertices: 3D coordinates of each vertex
257
+ - faces: triangles defined by vertex indices
258
+ - face_colors: RGBA colors for each face based on sim_grid values
259
+ - visual: trimesh.visual.ColorVisuals object storing the face colors
260
+
261
+ Examples
262
+ --------
263
+ Basic usage with Green View Index data:
264
+ >>> gvi = np.array([[0.5, 0.6], [0.4, 0.8]]) # GVI values
265
+ >>> dem = np.array([[10.0, 10.2], [9.8, 10.1]]) # Ground heights
266
+ >>> mesh = create_sim_surface_mesh(gvi, dem, meshsize=1.0, z_offset=1.5)
267
+
268
+ Custom color range and colormap:
269
+ >>> mesh = create_sim_surface_mesh(gvi, dem,
270
+ ... cmap_name='RdYlBu',
271
+ ... vmin=0.0, vmax=1.0)
272
+
273
+ Notes
274
+ -----
275
+ - The function automatically creates a matplotlib colorbar figure for visualization
276
+ - Both input grids are flipped vertically to match the voxel_array orientation
277
+ - Each grid cell is converted to two triangles for compatibility with 3D engines
278
+ - The mesh is positioned at dem_grid + z_offset to float above the terrain
279
+ - Face colors are interpolated from the colormap based on sim_grid values
280
+ """
281
+ # Flip arrays vertically
282
+ sim_grid_flipped = np.flipud(sim_grid)
283
+ dem_grid_flipped = np.flipud(dem_grid)
284
+
285
+ # Identify valid (non-NaN) values
286
+ valid_mask = ~np.isnan(sim_grid_flipped)
287
+ valid_values = sim_grid_flipped[valid_mask]
288
+ if valid_values.size == 0:
289
+ return None
290
+
291
+ # If vmin/vmax not provided, use actual min/max of the valid sim data
292
+ if vmin is None:
293
+ vmin = np.nanmin(valid_values)
294
+ if vmax is None:
295
+ vmax = np.nanmax(valid_values)
296
+
297
+ # Prepare the colormap and create colorbar
298
+ norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
299
+ scalar_map = cm.ScalarMappable(norm=norm, cmap=cmap_name)
300
+
301
+ # Create a figure just for the colorbar
302
+ fig, ax = plt.subplots(figsize=(6, 1))
303
+ plt.colorbar(scalar_map, cax=ax, orientation='horizontal')
304
+ plt.tight_layout()
305
+ plt.close()
306
+
307
+ vertices = []
308
+ faces = []
309
+ face_colors = []
310
+
311
+ vert_index = 0
312
+ nrows, ncols = sim_grid_flipped.shape
313
+
314
+ # Build a quad (two triangles) for each valid cell
315
+ for x in range(nrows):
316
+ for y in range(ncols):
317
+ val = sim_grid_flipped[x, y]
318
+ if np.isnan(val):
319
+ continue
320
+
321
+ z_base = meshsize * int(dem_grid_flipped[x, y] / meshsize + 1.5) + z_offset
322
+
323
+ # 4 corners in (x,y)*meshsize
324
+ v0 = [ x * meshsize, y * meshsize, z_base ]
325
+ v1 = [(x + 1) * meshsize, y * meshsize, z_base ]
326
+ v2 = [(x + 1) * meshsize, (y + 1) * meshsize, z_base ]
327
+ v3 = [ x * meshsize, (y + 1) * meshsize, z_base ]
328
+
329
+ vertices.extend([v0, v1, v2, v3])
330
+ faces.extend([
331
+ [vert_index, vert_index + 1, vert_index + 2],
332
+ [vert_index, vert_index + 2, vert_index + 3]
333
+ ])
334
+
335
+ # Get color from colormap
336
+ color_rgba = np.array(scalar_map.to_rgba(val)) # shape (4,)
337
+
338
+ # Each cell has 2 faces => add the color twice
339
+ face_colors.append(color_rgba)
340
+ face_colors.append(color_rgba)
341
+
342
+ vert_index += 4
343
+
344
+ if len(vertices) == 0:
345
+ return None
346
+
347
+ vertices = np.array(vertices, dtype=float)
348
+ faces = np.array(faces, dtype=int)
349
+ face_colors = np.array(face_colors, dtype=float)
350
+
351
+ mesh = trimesh.Trimesh(
352
+ vertices=vertices,
353
+ faces=faces,
354
+ face_colors=face_colors,
355
+ process=False # skip auto merge if you want to preserve quads
356
+ )
357
+
358
+ return mesh
359
+
360
+ def create_city_meshes(
361
+ voxel_array,
362
+ vox_dict,
363
+ meshsize=1.0,
364
+ include_classes=None,
365
+ exclude_classes=None,
366
+ ):
367
+ """
368
+ Create a collection of colored 3D meshes representing different city elements.
369
+
370
+ This function processes a voxelized city model and creates separate meshes for
371
+ different urban elements (buildings, trees, etc.), each with its own color.
372
+ The function preserves sharp edges and applies appropriate colors based on the
373
+ provided color dictionary.
374
+
375
+ Parameters
376
+ ----------
377
+ voxel_array : np.ndarray (3D)
378
+ 3D array representing the voxelized city model. Each voxel contains a class ID
379
+ that maps to an urban element type:
380
+ - 0: Void/air (automatically skipped)
381
+ - -2: Trees
382
+ - -3: Buildings
383
+ Other values can represent different urban elements as defined in vox_dict.
384
+
385
+ vox_dict : dict
386
+ Dictionary mapping class IDs to RGB colors. Each entry should be:
387
+ {class_id: [R, G, B]} where R, G, B are 0-255 integer values.
388
+ Example: {-3: [200, 200, 200], -2: [0, 255, 0]} for grey buildings and
389
+ green trees. The key 0 (air) is automatically excluded.
390
+
391
+ meshsize : float, default=1.0
392
+ Size of each voxel in meters, applied uniformly to x, y, and z dimensions.
393
+ Used to scale the output meshes to real-world coordinates.
394
+
395
+ Returns
396
+ -------
397
+ meshes : dict
398
+ Dictionary mapping class IDs to their corresponding trimesh.Trimesh objects.
399
+ Each mesh includes:
400
+ - vertices: 3D coordinates scaled by meshsize
401
+ - faces: triangulated faces preserving sharp edges
402
+ - face_colors: RGBA colors from vox_dict
403
+ - visual: trimesh.visual.ColorVisuals object storing the face colors
404
+
405
+ Classes with no voxels are automatically excluded from the output.
406
+
407
+ Examples
408
+ --------
409
+ Basic usage with buildings and trees:
410
+ >>> voxels = np.zeros((10, 10, 10))
411
+ >>> voxels[4:7, 4:7, 0:5] = -3 # Add a building
412
+ >>> voxels[2:4, 2:4, 0:3] = -2 # Add some trees
413
+ >>> colors = {
414
+ ... -3: [200, 200, 200], # Grey buildings
415
+ ... -2: [0, 255, 0] # Green trees
416
+ ... }
417
+ >>> meshes = create_city_meshes(voxels, colors, meshsize=1.0)
418
+
419
+ Notes
420
+ -----
421
+ - The function automatically skips class_id=0 (typically air/void)
422
+ - Each urban element type gets its own separate mesh for efficient rendering
423
+ - Colors are converted from RGB [0-255] to RGBA [0-1] format
424
+ - Sharp edges are preserved to maintain architectural features
425
+ - Empty classes (no voxels) are automatically excluded from the output
426
+ - Errors during mesh creation for a class are caught and reported
427
+ """
428
+ meshes = {}
429
+
430
+ # Convert RGB colors to hex for material properties
431
+ color_dict = {k: mcolors.rgb2hex([v[0]/255, v[1]/255, v[2]/255])
432
+ for k, v in vox_dict.items() if k != 0} # Exclude air
433
+
434
+ # Determine which classes to process
435
+ unique_classes = np.unique(voxel_array)
436
+
437
+ if include_classes is not None:
438
+ # Only keep classes explicitly requested (and present in the data)
439
+ class_iterable = [c for c in include_classes if c in unique_classes]
440
+ else:
441
+ class_iterable = list(unique_classes)
442
+
443
+ exclude_set = set(exclude_classes) if exclude_classes is not None else set()
444
+
445
+ # Create vertices and faces for each object class
446
+ for class_id in class_iterable:
447
+ if class_id == 0: # Skip air
448
+ continue
449
+
450
+ if class_id in exclude_set:
451
+ # Explicitly skipped (e.g., will be replaced with custom mesh)
452
+ continue
453
+
454
+ try:
455
+ mesh = create_voxel_mesh(voxel_array, class_id, meshsize=meshsize)
456
+
457
+ if mesh is None:
458
+ continue
459
+
460
+ # Convert hex color to RGBA
461
+ if class_id not in color_dict:
462
+ # Color not provided; skip silently for robustness
463
+ continue
464
+ rgb_color = np.array(mcolors.hex2color(color_dict[class_id]))
465
+ rgba_color = np.concatenate([rgb_color, [1.0]])
466
+
467
+ # Assign color to all faces
468
+ mesh.visual.face_colors = np.tile(rgba_color, (len(mesh.faces), 1))
469
+
470
+ meshes[class_id] = mesh
471
+
472
+ except ValueError as e:
473
+ print(f"Skipping class {class_id}: {e}")
474
+
475
+ return meshes
476
+
477
+ def export_meshes(meshes, output_directory, base_filename):
478
+ """
479
+ Export a collection of meshes to both OBJ (with MTL) and STL formats.
480
+
481
+ This function exports meshes in two ways:
482
+ 1. A single combined OBJ file with materials (and associated MTL file)
483
+ 2. Separate STL files for each mesh, named with their class IDs
484
+
485
+ Parameters
486
+ ----------
487
+ meshes : dict
488
+ Dictionary mapping class IDs to trimesh.Trimesh objects.
489
+ Each mesh should have:
490
+ - vertices: 3D coordinates
491
+ - faces: triangulated faces
492
+ - face_colors: RGBA colors (if using materials)
493
+
494
+ output_directory : str
495
+ Directory path where the output files will be saved.
496
+ Will be created if it doesn't exist.
497
+
498
+ base_filename : str
499
+ Base name for the output files (without extension).
500
+ Will be used to create:
501
+ - {base_filename}.obj : Combined mesh with materials
502
+ - {base_filename}.mtl : Material definitions for OBJ
503
+ - {base_filename}_{class_id}.stl : Individual STL files
504
+
505
+ Returns
506
+ -------
507
+ None
508
+ Files are written directly to the specified output directory.
509
+
510
+ Examples
511
+ --------
512
+ >>> meshes = {
513
+ ... -3: building_mesh, # Building mesh with grey color
514
+ ... -2: tree_mesh # Tree mesh with green color
515
+ ... }
516
+ >>> export_meshes(meshes, 'output/models', 'city_model')
517
+
518
+ This will create:
519
+ - output/models/city_model.obj
520
+ - output/models/city_model.mtl
521
+ - output/models/city_model_-3.stl
522
+ - output/models/city_model_-2.stl
523
+
524
+ Notes
525
+ -----
526
+ - OBJ/MTL format preserves colors and materials but is more complex
527
+ - STL format is simpler but doesn't support colors
528
+ - STL files are exported separately for each class for easier processing
529
+ - The OBJ file combines all meshes while preserving their materials
530
+ - File extensions are automatically added to the base filename
531
+ """
532
+ # Export combined mesh as OBJ with materials
533
+ combined_mesh = trimesh.util.concatenate(list(meshes.values()))
534
+ combined_mesh.export(f"{output_directory}/{base_filename}.obj")
535
+
536
+ # Export individual meshes as STL
537
+ for class_id, mesh in meshes.items():
538
+ # Convert class_id to a string for filename
539
+ mesh.export(f"{output_directory}/{base_filename}_{class_id}.stl")
540
+
541
+ def split_vertices_manual(mesh):
542
+ """
543
+ Split a mesh into independent faces by duplicating shared vertices.
544
+
545
+ This function imitates trimesh's split_vertices() functionality but ensures
546
+ complete face independence by giving each face its own copy of vertices.
547
+ This is particularly useful for rendering applications where smooth shading
548
+ between faces is undesirable, such as architectural visualization in Rhino.
549
+
550
+ Parameters
551
+ ----------
552
+ mesh : trimesh.Trimesh
553
+ Input mesh to split. Should have:
554
+ - vertices: array of vertex coordinates
555
+ - faces: array of vertex indices forming triangles
556
+ - visual: Optional ColorVisuals object with face colors
557
+
558
+ Returns
559
+ -------
560
+ out_mesh : trimesh.Trimesh
561
+ New mesh where each face is completely independent, with:
562
+ - Duplicated vertices for each face
563
+ - No vertex sharing between faces
564
+ - Preserved face colors if present in input
565
+ - Each face as a separate component
566
+
567
+ Examples
568
+ --------
569
+ Basic usage:
570
+ >>> vertices = np.array([[0,0,0], [1,0,0], [1,1,0], [0,1,0]])
571
+ >>> faces = np.array([[0,1,2], [0,2,3]]) # Two triangles sharing vertices
572
+ >>> mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
573
+ >>> split_mesh = split_vertices_manual(mesh)
574
+ >>> print(f"Original vertices: {len(mesh.vertices)}") # 4 vertices
575
+ >>> print(f"Split vertices: {len(split_mesh.vertices)}") # 6 vertices
576
+
577
+ With face colors:
578
+ >>> colors = np.array([[255,0,0,255], [0,255,0,255]]) # Red and green faces
579
+ >>> mesh.visual = trimesh.visual.ColorVisuals(mesh, face_colors=colors)
580
+ >>> split_mesh = split_vertices_manual(mesh) # Colors are preserved
581
+
582
+ Notes
583
+ -----
584
+ - Each output face has exactly 3 unique vertices
585
+ - Face colors are preserved in the output mesh
586
+ - Useful for:
587
+ - Preventing smooth shading artifacts
588
+ - Ensuring face color independence
589
+ - Preparing meshes for CAD software
590
+ - Creating sharp edges in architectural models
591
+ - Memory usage increases as vertices are duplicated
592
+ """
593
+ new_meshes = []
594
+
595
+ # For each face, build a small, one-face mesh
596
+ for face_idx, face in enumerate(mesh.faces):
597
+ face_coords = mesh.vertices[face]
598
+
599
+ # Create mini-mesh without colors first
600
+ mini_mesh = trimesh.Trimesh(
601
+ vertices=face_coords,
602
+ faces=[[0, 1, 2]],
603
+ process=False # skip merging/cleaning
604
+ )
605
+
606
+ # If the mesh has per-face colors, set the face color properly
607
+ if (mesh.visual.face_colors is not None
608
+ and len(mesh.visual.face_colors) == len(mesh.faces)):
609
+ # Create a visual object with the face color (for one face)
610
+ face_color = mesh.visual.face_colors[face_idx]
611
+ color_visual = trimesh.visual.ColorVisuals(
612
+ mesh=mini_mesh,
613
+ face_colors=np.array([face_color]), # One face, one color
614
+ vertex_colors=None
615
+ )
616
+ mini_mesh.visual = color_visual
617
+
618
+ new_meshes.append(mini_mesh)
619
+
620
+ # Concatenate all the single-face meshes
621
+ out_mesh = trimesh.util.concatenate(new_meshes)
622
+ return out_mesh
623
+
624
+ def save_obj_from_colored_mesh(meshes, output_path, base_filename, max_materials=None):
625
+ """
626
+ Memory-safe OBJ/MTL exporter.
627
+ - Streams vertices/faces to disk (no concatenate, no per-face mini-meshes).
628
+ - Uses face colors -> materials (no vertex splitting).
629
+ - Optional color quantization to reduce material count.
630
+ """
631
+ import os
632
+ import numpy as np
633
+
634
+ os.makedirs(output_path, exist_ok=True)
635
+ obj_path = os.path.join(output_path, f"{base_filename}.obj")
636
+ mtl_path = os.path.join(output_path, f"{base_filename}.mtl")
637
+
638
+ # --------------- helpers ---------------
639
+ def to_uint8_rgba(arr):
640
+ arr = np.asarray(arr)
641
+ if arr.dtype != np.uint8:
642
+ # Handle float [0..1] or int [0..255]
643
+ if arr.dtype.kind == 'f':
644
+ arr = np.clip(arr, 0.0, 1.0)
645
+ arr = (arr * 255.0 + 0.5).astype(np.uint8)
646
+ else:
647
+ arr = arr.astype(np.uint8)
648
+ if arr.shape[1] == 3:
649
+ alpha = np.full((arr.shape[0], 1), 255, dtype=np.uint8)
650
+ arr = np.concatenate([arr, alpha], axis=1)
651
+ return arr
652
+
653
+ # First pass: build material palette
654
+ # We avoid collecting all colors at once—scan per mesh and update a dict.
655
+ color_to_id = {}
656
+ ordered_colors = [] # list of RGBA uint8 tuples in material order
657
+
658
+ # Optional quantizer (lazy-init)
659
+ quantizer = None
660
+ if max_materials is not None:
661
+ try:
662
+ from sklearn.cluster import MiniBatchKMeans
663
+ quantizer = MiniBatchKMeans(n_clusters=max_materials, random_state=42, batch_size=8192)
664
+ # Partial-fit streaming pass over colors
665
+ for m in meshes.values():
666
+ fc = getattr(m.visual, "face_colors", None)
667
+ if fc is None:
668
+ continue
669
+ fc = to_uint8_rgba(fc)
670
+ if fc.size == 0:
671
+ continue
672
+ # Use only RGB for clustering
673
+ quantizer.partial_fit(fc[:, :3].astype(np.float32))
674
+ except ImportError:
675
+ raise ImportError("scikit-learn is required for color quantization. Install it with: pip install scikit-learn")
676
+
677
+ # Assign material ids during a second scan, but still streaming to avoid big unions
678
+ def get_material_id(rgba):
679
+ key = (int(rgba[0]), int(rgba[1]), int(rgba[2]), int(rgba[3]))
680
+ mid = color_to_id.get(key)
681
+ if mid is None:
682
+ mid = len(ordered_colors)
683
+ color_to_id[key] = mid
684
+ ordered_colors.append(key)
685
+ return mid
686
+
687
+ # 2nd pass if quantizing: we need color centroids
688
+ centers_u8 = None
689
+ if quantizer is not None:
690
+ centers = quantizer.cluster_centers_.astype(np.float32) # RGB float
691
+ centers = np.clip(centers, 0.0, 255.0).astype(np.uint8)
692
+ # Build a quick LUT fun
693
+ def quantize_rgb(rgb_u8):
694
+ # rgb_u8: (N,3) uint8 -> labels -> centers
695
+ labels = quantizer.predict(rgb_u8.astype(np.float32))
696
+ return centers[labels]
697
+ # We'll convert each mesh's face colors to quantized RGB on the fly
698
+ centers_u8 = centers
699
+
700
+ # Build materials palette by scanning once (still O(total faces) but tiny memory)
701
+ for m in meshes.values():
702
+ fc = getattr(m.visual, "face_colors", None)
703
+ if fc is None:
704
+ # No colors: assign default grey
705
+ rgba = np.array([[200,200,200,255]], dtype=np.uint8)
706
+ get_material_id(rgba[0])
707
+ continue
708
+ fc = to_uint8_rgba(fc)
709
+ if quantizer is not None:
710
+ q_rgb = quantize_rgb(fc[:, :3])
711
+ fc = np.concatenate([q_rgb, fc[:, 3:4]], axis=1)
712
+ # Iterate unique colors in this mesh to limit get_material_id calls
713
+ # but don't materialize huge sets; unique per mesh is fine.
714
+ uniq = np.unique(fc, axis=0)
715
+ for rgba in uniq:
716
+ get_material_id(rgba)
717
+
718
+ # Write MTL
719
+ with open(mtl_path, "w") as mtl:
720
+ for i, (r, g, b, a) in enumerate(ordered_colors):
721
+ mtl.write(f"newmtl material_{i}\n")
722
+ # Match viewport look: diffuse only, no specular. Many viewers assume sRGB.
723
+ kd_r, kd_g, kd_b = r/255.0, g/255.0, b/255.0
724
+ mtl.write(f"Kd {kd_r:.6f} {kd_g:.6f} {kd_b:.6f}\n")
725
+ # Ambient same as diffuse to avoid darkening in some viewers
726
+ mtl.write(f"Ka {kd_r:.6f} {kd_g:.6f} {kd_b:.6f}\n")
727
+ # No specular highlight
728
+ mtl.write("Ks 0.000000 0.000000 0.000000\n")
729
+ # Disable lighting model with specular; keep simple shading
730
+ mtl.write("illum 1\n")
731
+ # Alpha
732
+ mtl.write(f"d {a/255.0:.6f}\n\n")
733
+
734
+ # Stream OBJ
735
+ with open(obj_path, "w") as obj:
736
+ obj.write(f"mtllib {os.path.basename(mtl_path)}\n")
737
+
738
+ v_offset = 0 # running vertex index offset
739
+ # Reusable cache so we don't keep writing 'usemtl' for the same block unnecessarily
740
+ current_material = None
741
+
742
+ for class_id, m in meshes.items():
743
+ verts = np.asarray(m.vertices, dtype=np.float64)
744
+ faces = np.asarray(m.faces, dtype=np.int64)
745
+ if verts.size == 0 or faces.size == 0:
746
+ continue
747
+
748
+ # Write vertices
749
+ # (We do a single pass; writing text is the bottleneck, but memory-safe.)
750
+ for v in verts:
751
+ obj.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
752
+
753
+ # Prepare face colors (face-level)
754
+ fc = getattr(m.visual, "face_colors", None)
755
+ if fc is None or len(fc) != len(faces):
756
+ # default grey if missing or mismatched
757
+ fc = np.tile(np.array([200,200,200,255], dtype=np.uint8), (len(faces), 1))
758
+ else:
759
+ fc = to_uint8_rgba(fc)
760
+
761
+ if quantizer is not None:
762
+ q_rgb = quantize_rgb(fc[:, :3])
763
+ fc = np.concatenate([q_rgb, fc[:, 3:4]], axis=1)
764
+
765
+ # Group faces by material id and stream in order
766
+ # Build material id per face quickly
767
+ # Convert face colors to material ids
768
+ # (Avoid Python loops over faces more than once)
769
+ # Map unique colors in this mesh to material ids first:
770
+ uniq_colors, inv_idx = np.unique(fc, axis=0, return_inverse=True)
771
+ color_to_mid_local = {tuple(c.tolist()): get_material_id(c) for c in uniq_colors}
772
+ mids = np.fromiter(
773
+ (color_to_mid_local[tuple(c.tolist())] for c in uniq_colors[inv_idx]),
774
+ dtype=np.int64,
775
+ count=len(inv_idx)
776
+ )
777
+
778
+ # Write faces grouped by material, but preserve simple ordering
779
+ # Cheap approach: emit runs; switching material only when necessary
780
+ current_material = None
781
+ for i_face, face in enumerate(faces):
782
+ mid = int(mids[i_face])
783
+ if current_material != mid:
784
+ obj.write(f"usemtl material_{mid}\n")
785
+ current_material = mid
786
+ a, b, c = face + 1 + v_offset # OBJ is 1-based
787
+ obj.write(f"f {a} {b} {c}\n")
788
+
789
+ v_offset += len(verts)
790
+
791
791
  return obj_path, mtl_path