Coverage for src/arcade_collection/convert/convert_to_meshes.py: 100%
100 statements
« prev ^ index » next coverage.py v7.1.0, created at 2024-12-09 19:07 +0000
1from __future__ import annotations
3from enum import Enum
4from typing import TYPE_CHECKING
6import numpy as np
7from skimage import measure
9from arcade_collection.output.extract_tick_json import extract_tick_json
10from arcade_collection.output.get_location_voxels import get_location_voxels
12if TYPE_CHECKING:
13 import tarfile
15 import pandas as pd
17MAX_ARRAY_LEVEL = 7
18"""Maximum array level for conversion to meshes."""
21class MeshType(Enum):
22 """Mesh face types."""
24 DEFAULT = False
25 """Mesh with default faces."""
27 INVERTED = True
28 """Mesh with inverted faces."""
31def convert_to_meshes(
32 series_key: str,
33 locations_tar: tarfile.TarFile,
34 frame_spec: tuple[int, int, int],
35 regions: list[str],
36 box: tuple[int, int, int],
37 mesh_type: MeshType | dict[str, MeshType] = MeshType.DEFAULT,
38 group_size: int | None = None,
39 categories: pd.DataFrame | None = None,
40) -> list[tuple[int, int, str, str]]:
41 """
42 Convert data to mesh OBJ contents.
44 Parameters
45 ----------
46 series_key
47 Simulation series key.
48 locations_tar
49 Archive of location data.
50 frame_spec
51 Specification for mesh frames.
52 regions
53 List of regions.
54 box
55 Size of bounding box.
56 mesh_type
57 Mesh face type.
58 group_size
59 Number of objects in each group (if grouping meshes).
60 categories
61 Simulation data containing ID, FRAME, and CATEGORY.
63 Returns
64 -------
65 :
66 List of mesh frames, indices, regions, and OBJ contents.
67 """
69 frames = list(np.arange(*frame_spec))
70 meshes = []
72 length, width, height = box
73 groups = make_mesh_groups(categories, frames, group_size) if group_size is not None else None
75 for frame in frames:
76 locations = extract_tick_json(locations_tar, series_key, frame, "LOCATIONS")
78 for region in regions:
79 region_mesh_type = mesh_type[region] if isinstance(mesh_type, dict) else mesh_type
81 if groups is None:
82 for location in locations:
83 location_id = location["id"]
84 mesh = make_individual_mesh(
85 location, length, width, height, region, region_mesh_type
86 )
88 if mesh is None:
89 continue
91 meshes.append((frame, location_id, region, mesh))
92 else:
93 for index, group in groups[frame].items():
94 group_locations = [
95 location for location in locations if location["id"] in group
96 ]
97 mesh = make_combined_mesh(
98 group_locations, length, width, height, region, region_mesh_type
99 )
101 if mesh is None:
102 continue
104 meshes.append((frame, index, region, mesh))
106 return meshes
109def make_mesh_groups(
110 categories: pd.DataFrame, frames: list[int], group_size: int
111) -> dict[int, dict[int, list[int]]]:
112 """
113 Group objects based on group size and categories.
115 Parameters
116 ----------
117 categories
118 Simulation data containing ID, FRAME, and CATEGORY.
119 frames
120 List of frames.
121 group_size
122 Number of objects in each group.
124 Returns
125 -------
126 :
127 Map of frame to map of index to location ids.
128 """
130 groups: dict[int, dict[int, list[int]]] = {}
132 for frame in frames:
133 groups[frame] = {}
134 frame_categories = categories[categories["FRAME"] == frame]
135 index_offset = 0
137 for _, category_group in frame_categories.groupby("CATEGORY"):
138 ids = list(category_group["ID"].values)
139 group_ids = [ids[i : i + group_size] for i in range(0, len(ids), group_size)]
140 groups[frame].update({i + index_offset: group for i, group in enumerate(group_ids)})
142 index_offset = index_offset + len(group_ids)
144 return groups
147def make_individual_mesh(
148 location: dict,
149 length: int,
150 width: int,
151 height: int,
152 region: str,
153 mesh_type: MeshType = MeshType.DEFAULT,
154) -> str | None:
155 """
156 Create mesh containing a single object.
158 Parameters
159 ----------
160 location
161 Location object.
162 length
163 Bounding box length.
164 width
165 Bounding box width.
166 height
167 Bounding box height.
168 region
169 Region name.
170 mesh_type
171 Mesh face type.
173 Returns
174 -------
175 :
176 Single mesh OBJ file contents.
177 """
179 voxels = [
180 (x, width - y - 1, z)
181 for x, y, z in get_location_voxels(location, region if region != "DEFAULT" else None)
182 ]
184 if len(voxels) == 0:
185 return None
187 center = list(np.array(voxels).mean(axis=0))
188 array = make_mesh_array(voxels, length, width, height)
189 verts, faces, normals = make_mesh_geometry(array, center)
190 return make_mesh_file(verts, faces, normals, mesh_type)
193def make_combined_mesh(
194 locations: list[dict],
195 length: int,
196 width: int,
197 height: int,
198 region: str,
199 mesh_type: MeshType = MeshType.DEFAULT,
200) -> str | None:
201 """
202 Create mesh containing multiple objects.
204 Parameters
205 ----------
206 locations
207 List of location objects.
208 length
209 Bounding box length.
210 width
211 Bounding box width.
212 height
213 Bounding box height.
214 region
215 Region name.
216 mesh_type
217 Mesh face type.
219 Returns
220 -------
221 :
222 Combined mesh OBJ file contents.
223 """
225 meshes = []
226 offset = 0
228 for location in locations:
229 voxels = [
230 (x, width - y - 1, z)
231 for x, y, z in get_location_voxels(location, region if region != "DEFAULT" else None)
232 ]
234 if len(voxels) == 0:
235 continue
237 center = [length / 2, width / 2, height / 2]
238 array = make_mesh_array(voxels, length, width, height)
239 verts, faces, normals = make_mesh_geometry(array, center, offset)
240 mesh = make_mesh_file(verts, faces, normals, mesh_type)
242 meshes.append(mesh)
243 offset = offset + len(verts)
245 combined_mesh = "\n".join(meshes)
247 return combined_mesh if meshes else None
250def make_mesh_array(
251 voxels: list[tuple[int, int, int]], length: int, width: int, height: int
252) -> np.ndarray:
253 """
254 Generate array from list of voxels.
256 Given voxel locations are set to the max array level. The array is smoothed
257 such that all other locations are set to the number of max-level neighbors.
259 Parameters
260 ----------
261 voxels
262 List of voxels representing object.
263 length
264 Bounding box length.
265 width
266 Bounding box width.
267 height
268 Bounding box height.
270 Returns
271 -------
272 :
273 Array representing object.
274 """
276 # Create array.
277 array = np.zeros((length, width, height), dtype=np.uint8)
278 array[tuple(np.transpose(voxels))] = MAX_ARRAY_LEVEL
280 # Get set of zero neighbors for all voxels.
281 offsets = [(-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, -1), (0, 0, 1)]
282 neighbors = {
283 (x + i, y + j, z + k)
284 for x, y, z in voxels
285 for i, j, k in offsets
286 if array[x + i, y + j, z + k] == 0
287 }
289 # Remove invalid neighbors on borders.
290 neighbors = {
291 (x, y, z)
292 for x, y, z in neighbors
293 if x != 0 and y != 0 and z != 0 and x != length - 1 and y != width - 1 and z != height - 1
294 }
296 # Smooth array levels based on neighbor counts.
297 for x, y, z in neighbors:
298 array[x, y, z] = (
299 sum(array[x + i, y + j, z + k] == MAX_ARRAY_LEVEL for i, j, k in offsets) + 1
300 )
302 return array
305def make_mesh_geometry(
306 array: np.ndarray, center: list[float], offset: int = 0
307) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
308 """
309 Generate mesh from array.
311 Parameters
312 ----------
313 array
314 Array representing object.
315 center
316 Coordinate of object center.
317 offset
318 Offset for face indices.
320 Returns
321 -------
322 :
323 Arrays of mesh vertices, faces, and normals.
324 """
326 level = int(MAX_ARRAY_LEVEL / 2)
327 verts, faces, normals, _ = measure.marching_cubes(array, level=level, allow_degenerate=False)
329 # Center the vertices.
330 verts[:, 0] = verts[:, 0] - center[0]
331 verts[:, 1] = verts[:, 1] - center[1]
332 verts[:, 2] = verts[:, 2] - center[2]
334 # Adjust face indices.
335 faces = faces + 1 + offset
337 return verts, faces, normals
340def make_mesh_file(
341 verts: np.ndarray,
342 faces: np.ndarray,
343 normals: np.ndarray,
344 mesh_type: MeshType = MeshType.DEFAULT,
345) -> str:
346 """
347 Create mesh OBJ file contents from marching cubes output.
349 If
351 Parameters
352 ----------
353 verts
354 Array of mesh vertices.
355 faces
356 Array of mesh faces.
357 normals
358 Array of mesh normals.
359 mesh_type
360 Mesh face type.
362 Returns
363 -------
364 :
365 Mesh OBJ file.
366 """
368 mesh = ""
370 for item in verts:
371 mesh += f"v {item[0]} {item[1]} {item[2]}\n"
373 for item in normals:
374 mesh += f"vn {item[0]} {item[1]} {item[2]}\n"
376 for item in faces:
377 if mesh_type == MeshType.INVERTED:
378 mesh += f"f {item[0]}//{item[0]} {item[1]}//{item[1]} {item[2]}//{item[2]}\n"
379 else:
380 mesh += f"f {item[2]}//{item[2]} {item[1]}//{item[1]} {item[0]}//{item[0]}\n"
382 return mesh