voxcity 0.5.31__py3-none-any.whl → 0.6.1__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,759 +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
- # Add building IDs as metadata for buildings
191
- if class_id == -3 and building_id_grid is not None and building_ids:
192
- mesh.metadata = {'building_id': np.array(building_ids)}
193
-
194
- return mesh
195
-
196
- def create_sim_surface_mesh(sim_grid, dem_grid,
197
- meshsize=1.0, z_offset=1.5,
198
- cmap_name='viridis',
199
- vmin=None, vmax=None):
200
- """
201
- Create a colored planar surface mesh from simulation data, positioned above a DEM.
202
-
203
- This function generates a 3D visualization mesh for 2D simulation results (like
204
- Green View Index, solar radiation, etc.). The mesh is positioned above the Digital
205
- Elevation Model (DEM) by a specified offset, and colored according to the simulation
206
- values using a matplotlib colormap.
207
-
208
- Parameters
209
- ----------
210
- sim_grid : 2D np.ndarray
211
- 2D array of simulation values (e.g., Green View Index, solar radiation).
212
- NaN values in this grid will be skipped in the output mesh.
213
- The grid should be oriented with north at the top.
214
-
215
- dem_grid : 2D np.ndarray
216
- 2D array of ground elevations in meters. Must have the same shape as sim_grid.
217
- Used to position the visualization mesh at the correct height above terrain.
218
-
219
- meshsize : float, default=1.0
220
- Size of each cell in meters. Applied uniformly to x and y dimensions.
221
- Determines the resolution of the output mesh.
222
-
223
- z_offset : float, default=1.5
224
- Additional height offset in meters added to dem_grid elevations.
225
- Used to position the visualization above ground level for better visibility.
226
-
227
- cmap_name : str, default='viridis'
228
- Matplotlib colormap name used for coloring the mesh based on sim_grid values.
229
- Common options:
230
- - 'viridis': Default, perceptually uniform, colorblind-friendly
231
- - 'RdYlBu': Red-Yellow-Blue, good for diverging data
232
- - 'jet': Rainbow colormap (not recommended for scientific visualization)
233
-
234
- vmin : float, optional
235
- Minimum value for color mapping. If None, uses min of sim_grid (excluding NaN).
236
- Used to control the range of the colormap.
237
-
238
- vmax : float, optional
239
- Maximum value for color mapping. If None, uses max of sim_grid (excluding NaN).
240
- Used to control the range of the colormap.
241
-
242
- Returns
243
- -------
244
- mesh : trimesh.Trimesh or None
245
- A single mesh containing one colored square face (two triangles) per non-NaN cell.
246
- Returns None if there are no valid (non-NaN) cells in sim_grid.
247
-
248
- The mesh includes:
249
- - vertices: 3D coordinates of each vertex
250
- - faces: triangles defined by vertex indices
251
- - face_colors: RGBA colors for each face based on sim_grid values
252
- - visual: trimesh.visual.ColorVisuals object storing the face colors
253
-
254
- Examples
255
- --------
256
- Basic usage with Green View Index data:
257
- >>> gvi = np.array([[0.5, 0.6], [0.4, 0.8]]) # GVI values
258
- >>> dem = np.array([[10.0, 10.2], [9.8, 10.1]]) # Ground heights
259
- >>> mesh = create_sim_surface_mesh(gvi, dem, meshsize=1.0, z_offset=1.5)
260
-
261
- Custom color range and colormap:
262
- >>> mesh = create_sim_surface_mesh(gvi, dem,
263
- ... cmap_name='RdYlBu',
264
- ... vmin=0.0, vmax=1.0)
265
-
266
- Notes
267
- -----
268
- - The function automatically creates a matplotlib colorbar figure for visualization
269
- - Both input grids are flipped vertically to match the voxel_array orientation
270
- - Each grid cell is converted to two triangles for compatibility with 3D engines
271
- - The mesh is positioned at dem_grid + z_offset to float above the terrain
272
- - Face colors are interpolated from the colormap based on sim_grid values
273
- """
274
- # Flip arrays vertically
275
- sim_grid_flipped = np.flipud(sim_grid)
276
- dem_grid_flipped = np.flipud(dem_grid)
277
-
278
- # Identify valid (non-NaN) values
279
- valid_mask = ~np.isnan(sim_grid_flipped)
280
- valid_values = sim_grid_flipped[valid_mask]
281
- if valid_values.size == 0:
282
- return None
283
-
284
- # If vmin/vmax not provided, use actual min/max of the valid sim data
285
- if vmin is None:
286
- vmin = np.nanmin(valid_values)
287
- if vmax is None:
288
- vmax = np.nanmax(valid_values)
289
-
290
- # Prepare the colormap and create colorbar
291
- norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
292
- scalar_map = cm.ScalarMappable(norm=norm, cmap=cmap_name)
293
-
294
- # Create a figure just for the colorbar
295
- fig, ax = plt.subplots(figsize=(6, 1))
296
- plt.colorbar(scalar_map, cax=ax, orientation='horizontal')
297
- plt.tight_layout()
298
- plt.close()
299
-
300
- vertices = []
301
- faces = []
302
- face_colors = []
303
-
304
- vert_index = 0
305
- nrows, ncols = sim_grid_flipped.shape
306
-
307
- # Build a quad (two triangles) for each valid cell
308
- for x in range(nrows):
309
- for y in range(ncols):
310
- val = sim_grid_flipped[x, y]
311
- if np.isnan(val):
312
- continue
313
-
314
- z_base = meshsize * int(dem_grid_flipped[x, y] / meshsize + 1.5) + z_offset
315
-
316
- # 4 corners in (x,y)*meshsize
317
- v0 = [ x * meshsize, y * meshsize, z_base ]
318
- v1 = [(x + 1) * meshsize, y * meshsize, z_base ]
319
- v2 = [(x + 1) * meshsize, (y + 1) * meshsize, z_base ]
320
- v3 = [ x * meshsize, (y + 1) * meshsize, z_base ]
321
-
322
- vertices.extend([v0, v1, v2, v3])
323
- faces.extend([
324
- [vert_index, vert_index + 1, vert_index + 2],
325
- [vert_index, vert_index + 2, vert_index + 3]
326
- ])
327
-
328
- # Get color from colormap
329
- color_rgba = np.array(scalar_map.to_rgba(val)) # shape (4,)
330
-
331
- # Each cell has 2 faces => add the color twice
332
- face_colors.append(color_rgba)
333
- face_colors.append(color_rgba)
334
-
335
- vert_index += 4
336
-
337
- if len(vertices) == 0:
338
- return None
339
-
340
- vertices = np.array(vertices, dtype=float)
341
- faces = np.array(faces, dtype=int)
342
- face_colors = np.array(face_colors, dtype=float)
343
-
344
- mesh = trimesh.Trimesh(
345
- vertices=vertices,
346
- faces=faces,
347
- face_colors=face_colors,
348
- process=False # skip auto merge if you want to preserve quads
349
- )
350
-
351
- return mesh
352
-
353
- def create_city_meshes(voxel_array, vox_dict, meshsize=1.0):
354
- """
355
- Create a collection of colored 3D meshes representing different city elements.
356
-
357
- This function processes a voxelized city model and creates separate meshes for
358
- different urban elements (buildings, trees, etc.), each with its own color.
359
- The function preserves sharp edges and applies appropriate colors based on the
360
- provided color dictionary.
361
-
362
- Parameters
363
- ----------
364
- voxel_array : np.ndarray (3D)
365
- 3D array representing the voxelized city model. Each voxel contains a class ID
366
- that maps to an urban element type:
367
- - 0: Void/air (automatically skipped)
368
- - -2: Trees
369
- - -3: Buildings
370
- Other values can represent different urban elements as defined in vox_dict.
371
-
372
- vox_dict : dict
373
- Dictionary mapping class IDs to RGB colors. Each entry should be:
374
- {class_id: [R, G, B]} where R, G, B are 0-255 integer values.
375
- Example: {-3: [200, 200, 200], -2: [0, 255, 0]} for grey buildings and
376
- green trees. The key 0 (air) is automatically excluded.
377
-
378
- meshsize : float, default=1.0
379
- Size of each voxel in meters, applied uniformly to x, y, and z dimensions.
380
- Used to scale the output meshes to real-world coordinates.
381
-
382
- Returns
383
- -------
384
- meshes : dict
385
- Dictionary mapping class IDs to their corresponding trimesh.Trimesh objects.
386
- Each mesh includes:
387
- - vertices: 3D coordinates scaled by meshsize
388
- - faces: triangulated faces preserving sharp edges
389
- - face_colors: RGBA colors from vox_dict
390
- - visual: trimesh.visual.ColorVisuals object storing the face colors
391
-
392
- Classes with no voxels are automatically excluded from the output.
393
-
394
- Examples
395
- --------
396
- Basic usage with buildings and trees:
397
- >>> voxels = np.zeros((10, 10, 10))
398
- >>> voxels[4:7, 4:7, 0:5] = -3 # Add a building
399
- >>> voxels[2:4, 2:4, 0:3] = -2 # Add some trees
400
- >>> colors = {
401
- ... -3: [200, 200, 200], # Grey buildings
402
- ... -2: [0, 255, 0] # Green trees
403
- ... }
404
- >>> meshes = create_city_meshes(voxels, colors, meshsize=1.0)
405
-
406
- Notes
407
- -----
408
- - The function automatically skips class_id=0 (typically air/void)
409
- - Each urban element type gets its own separate mesh for efficient rendering
410
- - Colors are converted from RGB [0-255] to RGBA [0-1] format
411
- - Sharp edges are preserved to maintain architectural features
412
- - Empty classes (no voxels) are automatically excluded from the output
413
- - Errors during mesh creation for a class are caught and reported
414
- """
415
- meshes = {}
416
-
417
- # Convert RGB colors to hex for material properties
418
- color_dict = {k: mcolors.rgb2hex([v[0]/255, v[1]/255, v[2]/255])
419
- for k, v in vox_dict.items() if k != 0} # Exclude air
420
-
421
- # Create vertices and faces for each object class
422
- for class_id in np.unique(voxel_array):
423
- if class_id == 0: # Skip air
424
- continue
425
-
426
- try:
427
- mesh = create_voxel_mesh(voxel_array, class_id, meshsize=meshsize)
428
-
429
- if mesh is None:
430
- continue
431
-
432
- # Convert hex color to RGBA
433
- rgb_color = np.array(mcolors.hex2color(color_dict[class_id]))
434
- rgba_color = np.concatenate([rgb_color, [1.0]])
435
-
436
- # Assign color to all faces
437
- mesh.visual.face_colors = np.tile(rgba_color, (len(mesh.faces), 1))
438
-
439
- meshes[class_id] = mesh
440
-
441
- except ValueError as e:
442
- print(f"Skipping class {class_id}: {e}")
443
-
444
- return meshes
445
-
446
- def export_meshes(meshes, output_directory, base_filename):
447
- """
448
- Export a collection of meshes to both OBJ (with MTL) and STL formats.
449
-
450
- This function exports meshes in two ways:
451
- 1. A single combined OBJ file with materials (and associated MTL file)
452
- 2. Separate STL files for each mesh, named with their class IDs
453
-
454
- Parameters
455
- ----------
456
- meshes : dict
457
- Dictionary mapping class IDs to trimesh.Trimesh objects.
458
- Each mesh should have:
459
- - vertices: 3D coordinates
460
- - faces: triangulated faces
461
- - face_colors: RGBA colors (if using materials)
462
-
463
- output_directory : str
464
- Directory path where the output files will be saved.
465
- Will be created if it doesn't exist.
466
-
467
- base_filename : str
468
- Base name for the output files (without extension).
469
- Will be used to create:
470
- - {base_filename}.obj : Combined mesh with materials
471
- - {base_filename}.mtl : Material definitions for OBJ
472
- - {base_filename}_{class_id}.stl : Individual STL files
473
-
474
- Returns
475
- -------
476
- None
477
- Files are written directly to the specified output directory.
478
-
479
- Examples
480
- --------
481
- >>> meshes = {
482
- ... -3: building_mesh, # Building mesh with grey color
483
- ... -2: tree_mesh # Tree mesh with green color
484
- ... }
485
- >>> export_meshes(meshes, 'output/models', 'city_model')
486
-
487
- This will create:
488
- - output/models/city_model.obj
489
- - output/models/city_model.mtl
490
- - output/models/city_model_-3.stl
491
- - output/models/city_model_-2.stl
492
-
493
- Notes
494
- -----
495
- - OBJ/MTL format preserves colors and materials but is more complex
496
- - STL format is simpler but doesn't support colors
497
- - STL files are exported separately for each class for easier processing
498
- - The OBJ file combines all meshes while preserving their materials
499
- - File extensions are automatically added to the base filename
500
- """
501
- # Export combined mesh as OBJ with materials
502
- combined_mesh = trimesh.util.concatenate(list(meshes.values()))
503
- combined_mesh.export(f"{output_directory}/{base_filename}.obj")
504
-
505
- # Export individual meshes as STL
506
- for class_id, mesh in meshes.items():
507
- # Convert class_id to a string for filename
508
- mesh.export(f"{output_directory}/{base_filename}_{class_id}.stl")
509
-
510
- def split_vertices_manual(mesh):
511
- """
512
- Split a mesh into independent faces by duplicating shared vertices.
513
-
514
- This function imitates trimesh's split_vertices() functionality but ensures
515
- complete face independence by giving each face its own copy of vertices.
516
- This is particularly useful for rendering applications where smooth shading
517
- between faces is undesirable, such as architectural visualization in Rhino.
518
-
519
- Parameters
520
- ----------
521
- mesh : trimesh.Trimesh
522
- Input mesh to split. Should have:
523
- - vertices: array of vertex coordinates
524
- - faces: array of vertex indices forming triangles
525
- - visual: Optional ColorVisuals object with face colors
526
-
527
- Returns
528
- -------
529
- out_mesh : trimesh.Trimesh
530
- New mesh where each face is completely independent, with:
531
- - Duplicated vertices for each face
532
- - No vertex sharing between faces
533
- - Preserved face colors if present in input
534
- - Each face as a separate component
535
-
536
- Examples
537
- --------
538
- Basic usage:
539
- >>> vertices = np.array([[0,0,0], [1,0,0], [1,1,0], [0,1,0]])
540
- >>> faces = np.array([[0,1,2], [0,2,3]]) # Two triangles sharing vertices
541
- >>> mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
542
- >>> split_mesh = split_vertices_manual(mesh)
543
- >>> print(f"Original vertices: {len(mesh.vertices)}") # 4 vertices
544
- >>> print(f"Split vertices: {len(split_mesh.vertices)}") # 6 vertices
545
-
546
- With face colors:
547
- >>> colors = np.array([[255,0,0,255], [0,255,0,255]]) # Red and green faces
548
- >>> mesh.visual = trimesh.visual.ColorVisuals(mesh, face_colors=colors)
549
- >>> split_mesh = split_vertices_manual(mesh) # Colors are preserved
550
-
551
- Notes
552
- -----
553
- - Each output face has exactly 3 unique vertices
554
- - Face colors are preserved in the output mesh
555
- - Useful for:
556
- - Preventing smooth shading artifacts
557
- - Ensuring face color independence
558
- - Preparing meshes for CAD software
559
- - Creating sharp edges in architectural models
560
- - Memory usage increases as vertices are duplicated
561
- """
562
- new_meshes = []
563
-
564
- # For each face, build a small, one-face mesh
565
- for face_idx, face in enumerate(mesh.faces):
566
- face_coords = mesh.vertices[face]
567
-
568
- # Create mini-mesh without colors first
569
- mini_mesh = trimesh.Trimesh(
570
- vertices=face_coords,
571
- faces=[[0, 1, 2]],
572
- process=False # skip merging/cleaning
573
- )
574
-
575
- # If the mesh has per-face colors, set the face color properly
576
- if (mesh.visual.face_colors is not None
577
- and len(mesh.visual.face_colors) == len(mesh.faces)):
578
- # Create a visual object with the face color (for one face)
579
- face_color = mesh.visual.face_colors[face_idx]
580
- color_visual = trimesh.visual.ColorVisuals(
581
- mesh=mini_mesh,
582
- face_colors=np.array([face_color]), # One face, one color
583
- vertex_colors=None
584
- )
585
- mini_mesh.visual = color_visual
586
-
587
- new_meshes.append(mini_mesh)
588
-
589
- # Concatenate all the single-face meshes
590
- out_mesh = trimesh.util.concatenate(new_meshes)
591
- return out_mesh
592
-
593
- def save_obj_from_colored_mesh(meshes, output_path, base_filename, max_materials=None):
594
- """
595
- Save a collection of colored meshes as OBJ and MTL files with material support.
596
-
597
- This function exports colored meshes to the Wavefront OBJ format with an
598
- accompanying MTL file for material definitions. It handles the conversion of
599
- face colors to materials and ensures proper material assignment in the OBJ file.
600
- The function is particularly useful for preserving color information in
601
- architectural and urban visualization models.
602
-
603
- Parameters
604
- ----------
605
- meshes : dict
606
- Dictionary mapping class IDs to trimesh.Trimesh objects.
607
- Each mesh should have:
608
- - vertices: array of 3D coordinates
609
- - faces: array of vertex indices forming triangles
610
- - visual.face_colors: RGBA colors for each face
611
-
612
- output_path : str
613
- Directory path where to save the files.
614
- Will be created if it doesn't exist.
615
- Should be writable by the current user.
616
-
617
- base_filename : str
618
- Base name for the output files (without extension).
619
- Will be used to create:
620
- - {base_filename}.obj : The main geometry file
621
- - {base_filename}.mtl : The material definitions file
622
-
623
- max_materials : int, optional
624
- Maximum number of materials/colors to create. If specified and the number
625
- of unique colors exceeds this limit, colors will be quantized using
626
- k-means clustering to reduce them to the specified number.
627
- If None (default), all unique colors are preserved as separate materials.
628
-
629
- Returns
630
- -------
631
- tuple
632
- (obj_path, mtl_path) : Paths to the saved OBJ and MTL files.
633
- Both paths are absolute or relative depending on the input output_path.
634
-
635
- Examples
636
- --------
637
- Basic usage with multiple colored meshes:
638
- >>> building_mesh = trimesh.Trimesh(
639
- ... vertices=[[0,0,0], [1,0,0], [1,1,0]],
640
- ... faces=[[0,1,2]],
641
- ... face_colors=[[200,200,200,255]] # Grey color
642
- ... )
643
- >>> tree_mesh = trimesh.Trimesh(
644
- ... vertices=[[2,0,0], [3,0,0], [2.5,1,0]],
645
- ... faces=[[0,1,2]],
646
- ... face_colors=[[0,255,0,255]] # Green color
647
- ... )
648
- >>> meshes = {-3: building_mesh, -2: tree_mesh}
649
- >>> obj_path, mtl_path = save_obj_from_colored_mesh(
650
- ... meshes, 'output/models', 'city'
651
- ... )
652
-
653
- Limiting materials to 5:
654
- >>> obj_path, mtl_path = save_obj_from_colored_mesh(
655
- ... meshes, 'output/models', 'city', max_materials=5
656
- ... )
657
-
658
- Notes
659
- -----
660
- - Creates unique materials for each distinct face color
661
- - Material names are auto-generated as 'material_0', 'material_1', etc.
662
- - Handles both RGB and RGBA colors (alpha channel supported)
663
- - Colors are normalized from [0-255] to [0-1] range for MTL format
664
- - Vertices are written in OBJ's 1-based indexing format
665
- - Faces are grouped by material for efficient rendering
666
- - The MTL file is automatically referenced in the OBJ file
667
- - If max_materials is specified, k-means clustering is used for color quantization
668
- - Color quantization preserves the overall color distribution while reducing material count
669
- - Requires scikit-learn package when max_materials is specified (install with: pip install scikit-learn)
670
-
671
- File Format Details
672
- -----------------
673
- OBJ file structure:
674
- - mtllib reference to MTL file
675
- - All vertex coordinates (v)
676
- - Face definitions (f) grouped by material (usemtl)
677
-
678
- MTL file structure:
679
- - newmtl: Material name
680
- - Kd: Diffuse color (RGB)
681
- - d: Alpha/transparency
682
- """
683
- os.makedirs(output_path, exist_ok=True)
684
- obj_path = os.path.join(output_path, f"{base_filename}.obj")
685
- mtl_path = os.path.join(output_path, f"{base_filename}.mtl")
686
-
687
- # Combine all meshes
688
- combined_mesh = trimesh.util.concatenate(list(meshes.values()))
689
- combined_mesh = split_vertices_manual(combined_mesh)
690
-
691
- # Get face colors
692
- face_colors = combined_mesh.visual.face_colors
693
-
694
- # Apply color quantization if max_materials is specified
695
- if max_materials is not None:
696
- try:
697
- from sklearn.cluster import KMeans
698
- except ImportError:
699
- raise ImportError("scikit-learn is required for color quantization. "
700
- "Install it with: pip install scikit-learn")
701
-
702
- # Prepare colors for clustering (use only RGB, not alpha)
703
- colors_rgb = face_colors[:, :3].astype(float)
704
-
705
- # Perform k-means clustering
706
- kmeans = KMeans(n_clusters=max_materials, random_state=42, n_init=10)
707
- color_labels = kmeans.fit_predict(colors_rgb)
708
-
709
- # Get the cluster centers as the new colors
710
- quantized_colors_rgb = kmeans.cluster_centers_.astype(np.uint8)
711
-
712
- # Create new face colors with quantized RGB and original alpha
713
- quantized_face_colors = np.zeros_like(face_colors)
714
- # Assign each face to its nearest cluster center
715
- quantized_face_colors[:, :3] = quantized_colors_rgb[color_labels]
716
- quantized_face_colors[:, 3] = face_colors[:, 3] # Preserve original alpha
717
-
718
- # Update the mesh with quantized colors
719
- combined_mesh.visual.face_colors = quantized_face_colors
720
- face_colors = quantized_face_colors
721
-
722
- # Create unique materials for each unique face color
723
- unique_colors = np.unique(face_colors, axis=0)
724
-
725
- # Write MTL file
726
- with open(mtl_path, 'w') as mtl_file:
727
- for i, color in enumerate(unique_colors):
728
- material_name = f'material_{i}'
729
- mtl_file.write(f'newmtl {material_name}\n')
730
- # Convert RGBA to RGB float values
731
- rgb = color[:3].astype(float) / 255.0
732
- mtl_file.write(f'Kd {rgb[0]:.6f} {rgb[1]:.6f} {rgb[2]:.6f}\n')
733
- mtl_file.write(f'd {color[3]/255.0:.6f}\n\n') # Alpha value
734
-
735
- # Create material groups based on face colors
736
- color_to_material = {tuple(c): f'material_{i}' for i, c in enumerate(unique_colors)}
737
-
738
- # Write OBJ file
739
- with open(obj_path, 'w') as obj_file:
740
- obj_file.write(f'mtllib {os.path.basename(mtl_path)}\n')
741
-
742
- # Write vertices
743
- for vertex in combined_mesh.vertices:
744
- obj_file.write(f'v {vertex[0]:.6f} {vertex[1]:.6f} {vertex[2]:.6f}\n')
745
-
746
- # Write faces grouped by material
747
- current_material = None
748
- for face_idx, face in enumerate(combined_mesh.faces):
749
- face_color = tuple(face_colors[face_idx])
750
- material_name = color_to_material[face_color]
751
-
752
- if material_name != current_material:
753
- obj_file.write(f'usemtl {material_name}\n')
754
- current_material = material_name
755
-
756
- # OBJ indices are 1-based
757
- obj_file.write(f'f {face[0]+1} {face[1]+1} {face[2]+1}\n')
758
-
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
+
759
791
  return obj_path, mtl_path