"""Utility functions for working with meshes."""
import json
import shutil
from typing import List
import DracoPy
import meshio
from brainglobe_utils.general.system import get_num_processes
from cloudvolume.datasource.precomputed.mesh.multilod import (
MultiLevelPrecomputedMeshManifest,
)
from rich.progress import track
from scipy import sparse
from scipy.ndimage import (
binary_closing,
binary_fill_holes,
distance_transform_edt,
find_objects,
gaussian_filter,
)
from skimage.measure import marching_cubes
from brainglobe_atlasapi.structure_tree_util import (
get_structures_tree,
preorder_depth_first_search,
)
try:
from vedo import Mesh, Volume, write
except ModuleNotFoundError:
raise ModuleNotFoundError(
"Mesh generation with these utils requires vedo\n"
+ ' please install with "pip install vedo -U"'
)
import multiprocessing as mp
from pathlib import Path
import numpy as np
import zarr
from loguru import logger
from treelib import Node, Tree
from brainglobe_atlasapi.atlas_generation.volume_utils import (
create_masked_array,
)
# ----------------- #
# MESH CREATION #
# ----------------- #
def _smooth_constrained(mask, band_radius=4, max_iters=500, rel_tol=1e-6):
"""Smooth the 0.5 level-set of a 3D binary mask.
Implements "Surface Extraction from Binary Volumes with Higher-Order
Smoothness", Victor Lempitsky, CVPR 2010. Adapted from PyMCubes'
`smooth_constrained` (BSD-3-Clause, Copyright (c) 2012-2015 P. M.
Neila), with the matrix assembly vectorised.
Returns a float array whose 0 level-set is the smoothed surface. The
surface is constrained to stay within half a voxel of the mask boundary.
"""
mask = mask > 0
dist = np.where(
mask,
distance_transform_edt(mask) - 0.5,
0.5 - distance_transform_edt(~mask),
)
band = np.abs(dist) <= band_radius
# Second differences along each axis over band voxels; neighbours outside
# the band are treated as equal to the voxel itself.
n = np.count_nonzero(band)
idx = np.pad(np.full(band.shape, -1), 1, constant_values=-1)
idx[1:-1, 1:-1, 1:-1][band] = np.arange(n)
coords = np.argwhere(band) + 1
rows, cols, vals = [], [], []
for axis in range(3):
row = axis * n + np.arange(n)
diag = np.full(n, -2.0)
for step in (-1, 1):
neighbour_coords = coords.copy()
neighbour_coords[:, axis] += step
neighbour = idx[tuple(neighbour_coords.T)]
inside = neighbour >= 0
rows.append(row[inside])
cols.append(neighbour[inside])
vals.append(np.ones(inside.sum()))
diag += ~inside
rows.append(row)
cols.append(np.arange(n))
vals.append(diag)
q = sparse.csr_matrix(
(np.concatenate(vals), (np.concatenate(rows), np.concatenate(cols))),
shape=(3 * n, n),
)
q = (q.T @ q).tocsr()
# Weighted Jacobi iterations, clipped so voxels keep their inside/outside
# sign, stopping when energy improves by less than rel_tol per iteration
x = dist[band]
upper = np.where(x < 0, x, np.inf)
lower = np.where(x > 0, x, -np.inf)
upper[np.abs(upper) < 1] = 0
lower[np.abs(lower) < 1] = 0
inv_diag = 1.0 / q.diagonal()
off_diag = q - sparse.diags(q.diagonal())
tol = 1 - (1 - rel_tol) ** 10
energy = x @ (q @ x) / 2
for i in range(max_iters):
x = 0.5 * (-inv_diag * (off_diag @ x)) + 0.5 * x
x = np.clip(x, lower, upper)
if (i + 1) % 10 == 0:
energy, previous_energy = x @ (q @ x) / 2, energy
if (previous_energy - energy) / previous_energy < tol:
break
dist[band] = x
return dist
def _create_region_mesh(
meshes_dir_path: Path,
node: Node,
tree: Tree,
labels: list[int],
annotated_volume: np.ndarray | str | Path,
ROOT_ID: int,
closing_n_iters: int,
decimate_fraction: float,
smooth: bool,
verbosity: int = 0,
bbox: tuple[slice, ...] | None = None,
) -> None:
"""
Create and save an `.obj` mesh for a region and its descendants.
The mesh is generated from a binary mask built from `node.identifier` and
the identifiers of all child nodes in `tree`. Only labels present in
`annotated_volume` are considered. If no matching labels are found, or the
resulting mask is empty, no mesh is written.
`annotated_volume` may be provided as an in-memory NumPy array or as a
path to a zarr store, which will be opened in read mode.
For the root region (`node.identifier == ROOT_ID`), mesh extraction skips
the `closing_n_iters` argument. For all other regions, that parameter is
passed through to `extract_mesh_from_mask`.
Parameters
----------
meshes_dir_path : Path
Directory where mesh `.obj` files are written.
node
Tree node corresponding to the region whose mesh should be created.
tree
Structure hierarchy containing `node` and its descendants.
labels
Unique annotation labels present in `annotated_volume`, typically
`list(np.unique(annotated_volume))`.
annotated_volume : numpy.ndarray or str or Path
Annotation volume as a 3D array, or a path to a zarr store containing
the annotations.
ROOT_ID : int
Identifier of the root structure.
closing_n_iters : int
Number of morphological closing iterations to apply for non-root
regions during mesh extraction.
decimate_fraction : float
Fraction used to decimate the extracted mesh.
smooth : bool
Whether to smooth the extracted mesh.
verbosity : int, optional
Verbosity level used for debug output.
bbox : tuple of slice, optional
Region of `annotated_volume` containing the region's labels, padded
so morphological closing can't reach its edge. Only this crop is read
and meshed; the output is the same as meshing the full volume.
Raises
------
TypeError
If `annotated_volume` is neither a NumPy array nor a path to a zarr
store.
Returns
-------
None
Mesh data is written to disk when extraction succeeds.
"""
if verbosity > 0:
logger.debug(f"Creating mesh for region {node.identifier}")
# Avoid overwriting existing mesh
savepath = meshes_dir_path / f"{node.identifier}.obj"
# if savepath.exists():
# logger.debug(f"Mesh file save path exists already, skipping.")
# return
if not isinstance(annotated_volume, np.ndarray):
# If annotated_volume is a path to a zarr store, open it
if isinstance(annotated_volume, (str, Path)):
annotated_volume = zarr.open(annotated_volume, mode="r")
else:
raise TypeError(
"Argument annotated_volume should be a np.ndarray"
" or a path to a zarr store"
)
# Get labels for region and it's children
stree = tree.subtree(node.identifier)
ids = list(stree.nodes.keys())
# Keep only labels that are in the annotation volume
matched_labels = [i for i in ids if i in labels]
if (
not matched_labels
): # it fails if the region and all of its children are not in annotation
if verbosity > 0:
print(f"No labels found for {node.tag}")
return
else:
# Create mask and extract mesh
origin = None
if bbox is not None:
annotated_volume = annotated_volume[bbox]
origin = [s.start for s in bbox]
mask = create_masked_array(annotated_volume, ids)
if np.sum(mask) == 0:
print(f"Empty mask for {node.tag}")
else:
if node.identifier == ROOT_ID:
extract_mesh_from_mask(
mask,
obj_filepath=savepath,
smooth=smooth,
decimate_fraction=decimate_fraction,
origin=origin,
)
else:
extract_mesh_from_mask(
mask,
obj_filepath=savepath,
smooth=smooth,
closing_n_iters=closing_n_iters,
decimate_fraction=decimate_fraction,
origin=origin,
)
def _label_bounding_boxes(volume, labels, slab_depth=64):
"""
Return per-label inclusive (lo, hi) voxel bounds, each (len(labels), ndim).
`labels` must be sorted. Labels absent from `volume` keep lo > hi, and
voxel values not in `labels` are ignored. The volume is scanned in slabs
along axis 0 to avoid a full-size index array.
"""
lo = np.full((len(labels), volume.ndim), np.iinfo(np.int64).max)
hi = np.full((len(labels), volume.ndim), -1)
for z0 in range(0, volume.shape[0], slab_depth):
slab = volume[z0 : z0 + slab_depth]
# 1-based label index per voxel, 0 for values not in labels
idx = np.minimum(np.searchsorted(labels, slab), len(labels) - 1)
found = labels[idx] == slab
idx += 1
idx[~found] = 0
objects = find_objects(idx, max_label=len(labels))
for i, sl in enumerate(objects):
if sl is None:
continue
start = [s.start for s in sl]
stop = [s.stop - 1 for s in sl]
start[0] += z0
stop[0] += z0
lo[i] = np.minimum(lo[i], start)
hi[i] = np.maximum(hi[i], stop)
return lo, hi
[docs]
def create_region_mesh(args):
"""
Wrap _create_region_mesh which facilitates
multiprocessing.
"""
if not isinstance(args, (tuple, list)):
raise TypeError("args must be a tuple or list")
return _create_region_mesh(*args)
[docs]
def construct_meshes_from_annotation(
save_path: Path,
volume: np.ndarray,
structures_list,
closing_n_iters=2,
decimate_fraction=0,
smooth=False,
parallel: bool = True,
num_threads: int = -1,
verbosity: int = 0,
skip_structure_ids=None,
):
"""
Retrieve or construct atlas region meshes for a given annotation volume.
If an atlas is packaged with mesh files, reuse those. Otherwise, construct
the meshes using the existing volume and structure tree. Returns a
dictionary mapping structure IDs to their corresponding .obj mesh files.
Parameters
----------
save_path : Path
Path to the directory where new mesh files will be saved.
volume : np.ndarray
3D annotation volume.
structures_list : list
List of structure dictionaries containing id information.
smooth: bool
if True the surface mesh is smoothed
closing_n_iters: int
number of iterations of closing morphological operation.
set to None to avoid applying morphological operations
decimate_fraction: float in range [0, 1].
What fraction of the original number of vertices is to be kept.
EG .5 means that 50% of the vertices are kept,
the others are removed.
parallel: bool
If True, uses multiprocessing to speed up mesh creation
num_threads: int
Number of threads to use for parallel processing.
If -1, threads are set to the maximum number based on
available memory.
If > 0, uses that many threads.
verbosity: int
Level of verbosity for logging. 0 for no output, 1 for basic info.
skip_structure_ids: iterable of int or None
If provided, mesh generation for these structure IDs is skipped.
Returns
-------
dict
Dictionary of structure IDs and paths to their .obj mesh files.
"""
if num_threads == 0:
raise ValueError("Number of threads cannot be 0")
meshes_dir_path = save_path / "meshes"
meshes_dir_path.mkdir(exist_ok=True)
tree = get_structures_tree(structures_list)
labels = np.unique(volume).astype(np.uint32)
# Only used for parallel processing
ann_path = save_path / "temp_annotations.zarr"
for key, node in tree.nodes.items():
node.data = Region(key in labels)
# Mesh each region from its padded bounding box instead of the full
# volume. Padding by closing_n_iters + 1 keeps hole filling and closing
# identical to running them on the full volume.
label_lo, label_hi = _label_bounding_boxes(volume, labels)
label_index = {int(label): i for i, label in enumerate(labels)}
pad = (closing_n_iters or 0) + 1
volume_shape = volume.shape
def region_bbox(node):
# Root is closed with extract_mesh_from_mask's default iterations,
# not closing_n_iters, and spans almost the whole volume anyway.
if node.identifier == tree.root:
return None
idx = [
label_index[i]
for i in tree.subtree(node.identifier).nodes
if i in label_index
]
if not idx:
return None
lo = label_lo[idx].min(axis=0) - pad
hi = label_hi[idx].max(axis=0) + pad + 1
return tuple(
slice(max(int(a), 0), min(int(b), size))
for a, b, size in zip(lo, hi, volume_shape)
)
volume_size = volume.size
if parallel:
compressor = zarr.codecs.BloscCodec(
cname="zstd",
clevel=6,
shuffle="bitshuffle",
)
if ann_path.exists():
shutil.rmtree(ann_path)
ann_store = zarr.storage.LocalStore(ann_path)
zarr.create_array(
ann_store,
data=volume,
compressors=compressor,
)
ann_store.close()
volume = ann_path
root_id = tree.root
# Normalise skip set so filtering is a simple membership check
if skip_structure_ids is None:
skip_structure_ids = set()
elif not isinstance(skip_structure_ids, set):
skip_structure_ids = set(skip_structure_ids)
# Create a list of arguments for each region's mesh creation,
# filtering out structures that should be skipped upstream to
# avoid unnecessary inter-process communication.
args_list = [
(
meshes_dir_path,
node,
tree,
labels,
volume,
root_id,
closing_n_iters,
decimate_fraction,
smooth,
verbosity,
region_bbox(node),
)
for node in preorder_depth_first_search(tree)
if node.identifier not in skip_structure_ids
]
if parallel:
if num_threads == -1:
# Each thread uses ~ 7 times the number of voxels in the volume.
mem_per_thread = 7 * volume_size
num_threads = get_num_processes(
ram_needed_per_process=mem_per_thread,
n_max_processes=mp.cpu_count() - 1,
fraction_free_ram=0.05,
)
logger.info(f"Using {num_threads} threads for mesh creation")
with mp.Pool(num_threads) as pool:
for _ in track(
pool.imap(create_region_mesh, args_list),
total=len(args_list),
description="Creating meshes",
):
pass
shutil.rmtree(ann_path) # Clean up temporary annotations zarr store
else:
for args in track(
args_list, total=len(args_list), description="Creating meshes"
):
_create_region_mesh(*args)
meshes_dict = {}
structures_with_mesh = []
for s in structures_list:
mesh_path = meshes_dir_path / f'{s["id"]}.obj'
if not mesh_path.exists():
print(f"No mesh file exists for: {s}, ignoring it")
continue
if mesh_path.stat().st_size < 512:
print(f"obj file for {s} is too small, ignoring it.")
continue
structures_with_mesh.append(s)
meshes_dict[s["id"]] = mesh_path
print(
f"In the end, {len(structures_with_mesh)}"
" structures with mesh are kept",
)
return meshes_dict
[docs]
def write_mesh_info(
mesh_dir: Path,
vertex_quantization_bits: int = 16,
transform: List[int] = None,
lod_scale_multiplier: float = 1.0,
) -> dict:
"""
Write the mesh-directory metadata file.
Parameters
----------
mesh_dir : Path
Path to the mesh directory where the `info` file will be written.
vertex_quantization_bits : int, optional
Number of bits used for vertex quantization (default is 16).
transform : List[int], optional
A 4x3 transformation matrix in row-major order (default is identity).
lod_scale_multiplier : float, optional
Scale multiplier for level-of-detail (LOD) (default is 1.0).
Returns
-------
dict
The metadata dictionary that was written to the `info` file.
"""
if transform is None:
transform = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0]
info = {
"@type": "neuroglancer_multilod_draco",
"vertex_quantization_bits": int(vertex_quantization_bits),
"transform": [float(x) for x in transform],
"lod_scale_multiplier": float(lod_scale_multiplier),
}
(mesh_dir / "info").write_text(json.dumps(info, indent=2))
return info
[docs]
def write_mesh(
mesh: meshio.Mesh,
mesh_dir: Path,
segment_id: int,
vertex_quantization_bits: int = 16,
compression_level: int = 0,
):
"""Write one object as a single-LOD, single-fragment multi-res mesh.
Parameters
----------
mesh_dir : path to the mesh directory (gets `info`, `<id>`, `<id>.index`)
segment_id : int label of the object
vertices : (N, 3) float array of xyz coordinates ("points")
faces : (M, 3) int array of triangle vertex indices
vertex_quantization_bits : 10 or 16 (Neuroglancer only allows these two)
compression_level : 0-10 (higher = more compression)
"""
vertices = np.ascontiguousarray(mesh.points, dtype=np.float32)
faces = np.ascontiguousarray(mesh.cells[0].data, dtype=np.uint32)
# Define the single octree cell that holds the whole mesh.
# Draco quantizes ALL axes against one uniform range (a bounding *cube*).
bbox_min = vertices.min(axis=0)
bbox_max = vertices.max(axis=0)
qrange = float((bbox_max - bbox_min).max())
if qrange <= 0.0: # degenerate (single point / planar in all axes)
qrange = 1.0
grid_origin = bbox_min.astype(np.float32)
# Encode the one fragment
fragment = DracoPy.encode(
vertices,
faces,
quantization_bits=vertex_quantization_bits,
compression_level=compression_level,
quantization_range=qrange,
quantization_origin=grid_origin.tolist(),
)
manifest = MultiLevelPrecomputedMeshManifest(
segment_id=int(segment_id),
chunk_shape=np.array([qrange, qrange, qrange], dtype=np.float32),
grid_origin=grid_origin,
num_lods=1,
lod_scales=np.array([qrange], dtype=np.float32),
vertex_offsets=np.zeros((1, 3), dtype=np.float32),
num_fragments_per_lod=np.array([1], dtype=np.uint32),
fragment_positions=[np.zeros((1, 3), dtype=np.uint32)],
fragment_offsets=[np.array([len(fragment)], dtype=np.uint32)],
)
(mesh_dir / str(segment_id)).write_bytes(fragment)
(mesh_dir / f"{int(segment_id)}.index").write_bytes(manifest.to_binary())
[docs]
class Region(object):
"""
Class used to add metadata to treelib.Tree during atlas creation.
Using this means that you can then filter tree nodes depending on
whether they have a mesh/label.
"""
def __init__(self, has_label):
self.has_label = has_label
