# -*- coding: utf-8 -*- # Copyright (c) 2012-2015, P. M. Neila # All rights reserved. # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # * Redistributions of source code must retain the above copyright notice, this # list of conditions and the following disclaimer. # * Redistributions in binary form must reproduce the above copyright notice, # this list of conditions and the following disclaimer in the documentation # and/or other materials provided with the distribution. # * Neither the name of the copyright holder nor the names of its # contributors may be used to endorse or promote products derived from # this software without specific prior written permission. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ Utilities for smoothing the occ/sdf grids. """ import logging from typing import Tuple import numpy as np import torch import torch.nn.functional as F from scipy import ndimage as ndi from scipy import sparse __all__ = [ "smooth", "smooth_constrained", "smooth_gaussian", "signed_distance_function", "smooth_gpu", "smooth_constrained_gpu", "smooth_gaussian_gpu", "signed_distance_function_gpu", ] def _build_variable_indices(band: np.ndarray) -> np.ndarray: num_variables = np.count_nonzero(band) variable_indices = np.full(band.shape, -1, dtype=np.int_) variable_indices[band] = np.arange(num_variables) return variable_indices def _buildq3d(variable_indices: np.ndarray): """ Builds the filterq matrix for the given variables. """ num_variables = variable_indices.max() + 1 filterq = sparse.lil_matrix((3 * num_variables, num_variables)) # Pad variable_indices to simplify out-of-bounds accesses variable_indices = np.pad( variable_indices, [(0, 1), (0, 1), (0, 1)], mode="constant", constant_values=-1 ) coords = np.nonzero(variable_indices >= 0) for count, (i, j, k) in enumerate(zip(*coords)): assert variable_indices[i, j, k] == count filterq[3 * count, count] = -2 neighbor = variable_indices[i - 1, j, k] if neighbor >= 0: filterq[3 * count, neighbor] = 1 else: filterq[3 * count, count] += 1 neighbor = variable_indices[i + 1, j, k] if neighbor >= 0: filterq[3 * count, neighbor] = 1 else: filterq[3 * count, count] += 1 filterq[3 * count + 1, count] = -2 neighbor = variable_indices[i, j - 1, k] if neighbor >= 0: filterq[3 * count + 1, neighbor] = 1 else: filterq[3 * count + 1, count] += 1 neighbor = variable_indices[i, j + 1, k] if neighbor >= 0: filterq[3 * count + 1, neighbor] = 1 else: filterq[3 * count + 1, count] += 1 filterq[3 * count + 2, count] = -2 neighbor = variable_indices[i, j, k - 1] if neighbor >= 0: filterq[3 * count + 2, neighbor] = 1 else: filterq[3 * count + 2, count] += 1 neighbor = variable_indices[i, j, k + 1] if neighbor >= 0: filterq[3 * count + 2, neighbor] = 1 else: filterq[3 * count + 2, count] += 1 filterq = filterq.tocsr() return filterq.T.dot(filterq) def _buildq3d_gpu(variable_indices: torch.Tensor, chunk_size=10000): """ Builds the filterq matrix for the given variables on GPU, using chunking to reduce memory usage. """ device = variable_indices.device num_variables = variable_indices.max().item() + 1 # Pad variable_indices to simplify out-of-bounds accesses variable_indices = torch.nn.functional.pad( variable_indices, (0, 1, 0, 1, 0, 1), mode="constant", value=-1 ) coords = torch.nonzero(variable_indices >= 0) i, j, k = coords[:, 0], coords[:, 1], coords[:, 2] # Function to process a chunk of data def process_chunk(start, end): row_indices = [] col_indices = [] values = [] for axis in range(3): row_indices.append(3 * torch.arange(start, end, device=device) + axis) col_indices.append( variable_indices[i[start:end], j[start:end], k[start:end]] ) values.append(torch.full((end - start,), -2, device=device)) for offset in [-1, 1]: if axis == 0: neighbor = variable_indices[ i[start:end] + offset, j[start:end], k[start:end] ] elif axis == 1: neighbor = variable_indices[ i[start:end], j[start:end] + offset, k[start:end] ] else: neighbor = variable_indices[ i[start:end], j[start:end], k[start:end] + offset ] mask = neighbor >= 0 row_indices.append( 3 * torch.arange(start, end, device=device)[mask] + axis ) col_indices.append(neighbor[mask]) values.append(torch.ones(mask.sum(), device=device)) # Add 1 to the diagonal for out-of-bounds neighbors row_indices.append( 3 * torch.arange(start, end, device=device)[~mask] + axis ) col_indices.append( variable_indices[i[start:end], j[start:end], k[start:end]][~mask] ) values.append(torch.ones((~mask).sum(), device=device)) return torch.cat(row_indices), torch.cat(col_indices), torch.cat(values) # Process data in chunks all_row_indices = [] all_col_indices = [] all_values = [] for start in range(0, coords.shape[0], chunk_size): end = min(start + chunk_size, coords.shape[0]) row_indices, col_indices, values = process_chunk(start, end) all_row_indices.append(row_indices) all_col_indices.append(col_indices) all_values.append(values) # Concatenate all chunks row_indices = torch.cat(all_row_indices) col_indices = torch.cat(all_col_indices) values = torch.cat(all_values) # Create sparse tensor indices = torch.stack([row_indices, col_indices]) filterq = torch.sparse_coo_tensor( indices, values, (3 * num_variables, num_variables) ) # Compute filterq.T @ filterq return torch.sparse.mm(filterq.t(), filterq) # Usage example: # variable_indices = torch.tensor(...).cuda() # Your input tensor on GPU # result = _buildq3d_gpu(variable_indices) def _buildq2d(variable_indices: np.ndarray): """ Builds the filterq matrix for the given variables. Version for 2 dimensions. """ num_variables = variable_indices.max() + 1 filterq = sparse.lil_matrix((3 * num_variables, num_variables)) # Pad variable_indices to simplify out-of-bounds accesses variable_indices = np.pad( variable_indices, [(0, 1), (0, 1)], mode="constant", constant_values=-1 ) coords = np.nonzero(variable_indices >= 0) for count, (i, j) in enumerate(zip(*coords)): assert variable_indices[i, j] == count filterq[2 * count, count] = -2 neighbor = variable_indices[i - 1, j] if neighbor >= 0: filterq[2 * count, neighbor] = 1 else: filterq[2 * count, count] += 1 neighbor = variable_indices[i + 1, j] if neighbor >= 0: filterq[2 * count, neighbor] = 1 else: filterq[2 * count, count] += 1 filterq[2 * count + 1, count] = -2 neighbor = variable_indices[i, j - 1] if neighbor >= 0: filterq[2 * count + 1, neighbor] = 1 else: filterq[2 * count + 1, count] += 1 neighbor = variable_indices[i, j + 1] if neighbor >= 0: filterq[2 * count + 1, neighbor] = 1 else: filterq[2 * count + 1, count] += 1 filterq = filterq.tocsr() return filterq.T.dot(filterq) def _jacobi( filterq, x0: np.ndarray, lower_bound: np.ndarray, upper_bound: np.ndarray, max_iters: int = 10, rel_tol: float = 1e-6, weight: float = 0.5, ): """Jacobi method with constraints.""" jacobi_r = sparse.lil_matrix(filterq) shp = jacobi_r.shape jacobi_d = 1.0 / filterq.diagonal() jacobi_r.setdiag((0,) * shp[0]) jacobi_r = jacobi_r.tocsr() x = x0 # We check the stopping criterion each 10 iterations check_each = 10 cum_rel_tol = 1 - (1 - rel_tol) ** check_each energy_now = np.dot(x, filterq.dot(x)) / 2 logging.info("Energy at iter %d: %.6g", 0, energy_now) for i in range(max_iters): x_1 = -jacobi_d * jacobi_r.dot(x) x = weight * x_1 + (1 - weight) * x # Constraints. x = np.maximum(x, lower_bound) x = np.minimum(x, upper_bound) # Stopping criterion if (i + 1) % check_each == 0: # Update energy energy_before = energy_now energy_now = np.dot(x, filterq.dot(x)) / 2 logging.info("Energy at iter %d: %.6g", i + 1, energy_now) # Check stopping criterion cum_rel_improvement = (energy_before - energy_now) / energy_before if cum_rel_improvement < cum_rel_tol: break return x def signed_distance_function( levelset: np.ndarray, band_radius: int ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Return the distance to the 0.5 levelset of a function, the mask of the border (i.e., the nearest cells to the 0.5 level-set) and the mask of the band (i.e., the cells of the function whose distance to the 0.5 level-set is less of equal to `band_radius`). """ binary_array = np.where(levelset > 0, True, False) # Compute the band and the border. dist_func = ndi.distance_transform_edt distance = np.where( binary_array, dist_func(binary_array) - 0.5, -dist_func(~binary_array) + 0.5 ) border = np.abs(distance) < 1 band = np.abs(distance) <= band_radius return distance, border, band def signed_distance_function_iso0( levelset: np.ndarray, band_radius: int ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Return the distance to the 0 levelset of a function, the mask of the border (i.e., the nearest cells to the 0 level-set) and the mask of the band (i.e., the cells of the function whose distance to the 0 level-set is less of equal to `band_radius`). """ binary_array = levelset > 0 # Compute the band and the border. dist_func = ndi.distance_transform_edt distance = np.where( binary_array, dist_func(binary_array), -dist_func(~binary_array) ) border = np.zeros_like(levelset, dtype=bool) border[:-1, :, :] |= levelset[:-1, :, :] * levelset[1:, :, :] <= 0 border[:, :-1, :] |= levelset[:, :-1, :] * levelset[:, 1:, :] <= 0 border[:, :, :-1] |= levelset[:, :, :-1] * levelset[:, :, 1:] <= 0 band = np.abs(distance) <= band_radius return distance, border, band def signed_distance_function_gpu(levelset: torch.Tensor, band_radius: int): binary_array = (levelset > 0).float() # Compute distance transform dist_pos = ( F.max_pool3d( -binary_array.unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1 ) .squeeze(0) .squeeze(0) + binary_array ) dist_neg = F.max_pool3d( (binary_array - 1).unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1 ).squeeze(0).squeeze(0) + (1 - binary_array) distance = torch.where(binary_array > 0, dist_pos - 0.5, -dist_neg + 0.5) # breakpoint() # Use levelset as distance directly # distance = levelset # print(distance.shape) # Compute border and band border = torch.abs(distance) < 1 band = torch.abs(distance) <= band_radius return distance, border, band def smooth_constrained( binary_array: np.ndarray, band_radius: int = 4, max_iters: int = 250, rel_tol: float = 1e-6, ) -> np.ndarray: """ Implementation of the smoothing method from "Surface Extraction from Binary Volumes with Higher-Order Smoothness" Victor Lempitsky, CVPR10 """ # # Compute the distance map, the border and the band. logging.info("Computing distance transform...") # distance, _, band = signed_distance_function(binary_array, band_radius) binary_array_gpu = torch.from_numpy(binary_array).cuda() distance, _, band = signed_distance_function_gpu(binary_array_gpu, band_radius) distance = distance.cpu().numpy() band = band.cpu().numpy() variable_indices = _build_variable_indices(band) # Compute filterq. logging.info("Building matrix filterq...") if binary_array.ndim == 3: filterq = _buildq3d(variable_indices) # variable_indices_gpu = torch.from_numpy(variable_indices).cuda() # filterq_gpu = _buildq3d_gpu(variable_indices_gpu) # filterq = filterq_gpu.cpu().numpy() elif binary_array.ndim == 2: filterq = _buildq2d(variable_indices) else: raise ValueError("binary_array.ndim not in [2, 3]") # Initialize the variables. res = np.asarray(distance, dtype=np.double) x = res[band] upper_bound = np.where(x < 0, x, np.inf) lower_bound = np.where(x > 0, x, -np.inf) upper_bound[np.abs(upper_bound) < 1] = 0 lower_bound[np.abs(lower_bound) < 1] = 0 # Solve. logging.info("Minimizing energy...") x = _jacobi( filterq=filterq, x0=x, lower_bound=lower_bound, upper_bound=upper_bound, max_iters=max_iters, rel_tol=rel_tol, ) res[band] = x return res def total_variation_denoising(x, weight=0.1, num_iterations=5, eps=1e-8): diff_x = torch.diff(x, dim=0, prepend=x[:1]) diff_y = torch.diff(x, dim=1, prepend=x[:, :1]) diff_z = torch.diff(x, dim=2, prepend=x[:, :, :1]) norm = torch.sqrt(diff_x**2 + diff_y**2 + diff_z**2 + eps) div_x = torch.diff(diff_x / norm, dim=0, append=diff_x[-1:] / norm[-1:]) div_y = torch.diff(diff_y / norm, dim=1, append=diff_y[:, -1:] / norm[:, -1:]) div_z = torch.diff(diff_z / norm, dim=2, append=diff_z[:, :, -1:] / norm[:, :, -1:]) return x - weight * (div_x + div_y + div_z) def smooth_constrained_gpu( binary_array: torch.Tensor, band_radius: int = 4, max_iters: int = 250, rel_tol: float = 1e-4, ): distance, _, band = signed_distance_function_gpu(binary_array, band_radius) # Initialize variables x = distance[band] upper_bound = torch.where(x < 0, x, torch.tensor(float("inf"), device=x.device)) lower_bound = torch.where(x > 0, x, torch.tensor(float("-inf"), device=x.device)) upper_bound[torch.abs(upper_bound) < 1] = 0 lower_bound[torch.abs(lower_bound) < 1] = 0 # Define the 3D Laplacian kernel laplacian_kernel = torch.tensor( [ [ [ [[0, 1, 0], [1, -6, 1], [0, 1, 0]], [[1, 0, 1], [0, 0, 0], [1, 0, 1]], [[0, 1, 0], [1, 0, 1], [0, 1, 0]], ] ] ], device=x.device, ).float() laplacian_kernel = laplacian_kernel / laplacian_kernel.abs().sum() breakpoint() # Simplified Jacobi iteration for i in range(max_iters): # Reshape x to 5D tensor (batch, channel, depth, height, width) x_5d = x.view(1, 1, *band.shape) x_3d = x.view(*band.shape) # Apply 3D convolution laplacian = F.conv3d(x_5d, laplacian_kernel, padding=1) # Reshape back to original dimensions laplacian = laplacian.view(x.shape) # Use a small relaxation factor to improve stability relaxation_factor = 0.1 tv_weight = 0.1 # x_new = x + relaxation_factor * laplacian x_new = total_variation_denoising(x_3d, weight=tv_weight) # Print laplacian min and max # print(f"Laplacian min: {laplacian.min().item():.4f}, max: {laplacian.max().item():.4f}") # Apply constraints # Reshape x_new to match the dimensions of lower_bound and upper_bound x_new = x_new.view(x.shape) x_new = torch.clamp(x_new, min=lower_bound, max=upper_bound) # Check for convergence diff_norm = torch.norm(x_new - x) print(diff_norm) x_norm = torch.norm(x) if x_norm > 1e-8: # Avoid division by very small numbers relative_change = diff_norm / x_norm if relative_change < rel_tol: break elif diff_norm < rel_tol: # If x_norm is very small, check absolute change break x = x_new # Check for NaN and break if found, also check for inf if torch.isnan(x).any() or torch.isinf(x).any(): print(f"NaN or Inf detected at iteration {i}") breakpoint() break result = distance.clone() result[band] = x return result def smooth_gaussian(binary_array: np.ndarray, sigma: float = 3) -> np.ndarray: vol = np.float_(binary_array) - 0.5 return ndi.gaussian_filter(vol, sigma=sigma) def smooth_gaussian_gpu(binary_array: torch.Tensor, sigma: float = 3): # vol = binary_array.float() vol = binary_array kernel_size = int(2 * sigma + 1) kernel = torch.ones( 1, 1, kernel_size, kernel_size, kernel_size, device=binary_array.device, dtype=vol.dtype, ) / (kernel_size**3) return F.conv3d( vol.unsqueeze(0).unsqueeze(0), kernel, padding=kernel_size // 2 ).squeeze() def smooth(binary_array: np.ndarray, method: str = "auto", **kwargs) -> np.ndarray: """ Smooths the 0.5 level-set of a binary array. Returns a floating-point array with a smoothed version of the original level-set in the 0 isovalue. This function can apply two different methods: - A constrained smoothing method which preserves details and fine structures, but it is slow and requires a large amount of memory. This method is recommended when the input array is small (smaller than (500, 500, 500)). - A Gaussian filter applied over the binary array. This method is fast, but not very precise, as it can destroy fine details. It is only recommended when the input array is large and the 0.5 level-set does not contain thin structures. Parameters ---------- binary_array : ndarray Input binary array with the 0.5 level-set to smooth. method : str, one of ['auto', 'gaussian', 'constrained'] Smoothing method. If 'auto' is given, the method will be automatically chosen based on the size of `binary_array`. Parameters for 'gaussian' ------------------------- sigma : float Size of the Gaussian filter (default 3). Parameters for 'constrained' ---------------------------- max_iters : positive integer Number of iterations of the constrained optimization method (default 250). rel_tol: float Relative tolerance as a stopping criterion (default 1e-6). Output ------ res : ndarray Floating-point array with a smoothed 0 level-set. """ binary_array = np.asarray(binary_array) if method == "auto": if binary_array.size > 512**3: method = "gaussian" else: method = "constrained" if method == "gaussian": return smooth_gaussian(binary_array, **kwargs) if method == "constrained": return smooth_constrained(binary_array, **kwargs) raise ValueError("Unknown method '{}'".format(method)) def smooth_gpu(binary_array: torch.Tensor, method: str = "auto", **kwargs): if method == "auto": method = "gaussian" if binary_array.numel() > 512**3 else "constrained" if method == "gaussian": return smooth_gaussian_gpu(binary_array, **kwargs) elif method == "constrained": return smooth_constrained_gpu(binary_array, **kwargs) else: raise ValueError(f"Unknown method '{method}'")