tohid4n/PartCrafter

1# -*- coding: utf-8 -*-2 3# Copyright (c) 2012-2015, P. M. Neila4# All rights reserved.5 6# Redistribution and use in source and binary forms, with or without7# modification, are permitted provided that the following conditions are met:8 9# * Redistributions of source code must retain the above copyright notice, this10#   list of conditions and the following disclaimer.11 12# * Redistributions in binary form must reproduce the above copyright notice,13#   this list of conditions and the following disclaimer in the documentation14#   and/or other materials provided with the distribution.15 16# * Neither the name of the copyright holder nor the names of its17#   contributors may be used to endorse or promote products derived from18#   this software without specific prior written permission.19 20# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"21# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE22# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE23# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE24# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL25# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR26# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER27# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,28# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE29# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.30 31"""32Utilities for smoothing the occ/sdf grids.33"""34 35import logging36from typing import Tuple37 38import numpy as np39import torch40import torch.nn.functional as F41from scipy import ndimage as ndi42from scipy import sparse43 44__all__ = [45    "smooth",46    "smooth_constrained",47    "smooth_gaussian",48    "signed_distance_function",49    "smooth_gpu",50    "smooth_constrained_gpu",51    "smooth_gaussian_gpu",52    "signed_distance_function_gpu",53]54 55 56def _build_variable_indices(band: np.ndarray) -> np.ndarray:57    num_variables = np.count_nonzero(band)58    variable_indices = np.full(band.shape, -1, dtype=np.int_)59    variable_indices[band] = np.arange(num_variables)60    return variable_indices61 62 63def _buildq3d(variable_indices: np.ndarray):64    """65    Builds the filterq matrix for the given variables.66    """67 68    num_variables = variable_indices.max() + 169    filterq = sparse.lil_matrix((3 * num_variables, num_variables))70 71    # Pad variable_indices to simplify out-of-bounds accesses72    variable_indices = np.pad(73        variable_indices, [(0, 1), (0, 1), (0, 1)], mode="constant", constant_values=-174    )75 76    coords = np.nonzero(variable_indices >= 0)77    for count, (i, j, k) in enumerate(zip(*coords)):78 79        assert variable_indices[i, j, k] == count80 81        filterq[3 * count, count] = -282        neighbor = variable_indices[i - 1, j, k]83        if neighbor >= 0:84            filterq[3 * count, neighbor] = 185        else:86            filterq[3 * count, count] += 187 88        neighbor = variable_indices[i + 1, j, k]89        if neighbor >= 0:90            filterq[3 * count, neighbor] = 191        else:92            filterq[3 * count, count] += 193 94        filterq[3 * count + 1, count] = -295        neighbor = variable_indices[i, j - 1, k]96        if neighbor >= 0:97            filterq[3 * count + 1, neighbor] = 198        else:99            filterq[3 * count + 1, count] += 1100 101        neighbor = variable_indices[i, j + 1, k]102        if neighbor >= 0:103            filterq[3 * count + 1, neighbor] = 1104        else:105            filterq[3 * count + 1, count] += 1106 107        filterq[3 * count + 2, count] = -2108        neighbor = variable_indices[i, j, k - 1]109        if neighbor >= 0:110            filterq[3 * count + 2, neighbor] = 1111        else:112            filterq[3 * count + 2, count] += 1113 114        neighbor = variable_indices[i, j, k + 1]115        if neighbor >= 0:116            filterq[3 * count + 2, neighbor] = 1117        else:118            filterq[3 * count + 2, count] += 1119 120    filterq = filterq.tocsr()121    return filterq.T.dot(filterq)122 123 124def _buildq3d_gpu(variable_indices: torch.Tensor, chunk_size=10000):125    """126    Builds the filterq matrix for the given variables on GPU, using chunking to reduce memory usage.127    """128    device = variable_indices.device129    num_variables = variable_indices.max().item() + 1130 131    # Pad variable_indices to simplify out-of-bounds accesses132    variable_indices = torch.nn.functional.pad(133        variable_indices, (0, 1, 0, 1, 0, 1), mode="constant", value=-1134    )135 136    coords = torch.nonzero(variable_indices >= 0)137    i, j, k = coords[:, 0], coords[:, 1], coords[:, 2]138 139    # Function to process a chunk of data140    def process_chunk(start, end):141        row_indices = []142        col_indices = []143        values = []144 145        for axis in range(3):146            row_indices.append(3 * torch.arange(start, end, device=device) + axis)147            col_indices.append(148                variable_indices[i[start:end], j[start:end], k[start:end]]149            )150            values.append(torch.full((end - start,), -2, device=device))151 152            for offset in [-1, 1]:153                if axis == 0:154                    neighbor = variable_indices[155                        i[start:end] + offset, j[start:end], k[start:end]156                    ]157                elif axis == 1:158                    neighbor = variable_indices[159                        i[start:end], j[start:end] + offset, k[start:end]160                    ]161                else:162                    neighbor = variable_indices[163                        i[start:end], j[start:end], k[start:end] + offset164                    ]165 166                mask = neighbor >= 0167                row_indices.append(168                    3 * torch.arange(start, end, device=device)[mask] + axis169                )170                col_indices.append(neighbor[mask])171                values.append(torch.ones(mask.sum(), device=device))172 173                # Add 1 to the diagonal for out-of-bounds neighbors174                row_indices.append(175                    3 * torch.arange(start, end, device=device)[~mask] + axis176                )177                col_indices.append(178                    variable_indices[i[start:end], j[start:end], k[start:end]][~mask]179                )180                values.append(torch.ones((~mask).sum(), device=device))181 182        return torch.cat(row_indices), torch.cat(col_indices), torch.cat(values)183 184    # Process data in chunks185    all_row_indices = []186    all_col_indices = []187    all_values = []188 189    for start in range(0, coords.shape[0], chunk_size):190        end = min(start + chunk_size, coords.shape[0])191        row_indices, col_indices, values = process_chunk(start, end)192        all_row_indices.append(row_indices)193        all_col_indices.append(col_indices)194        all_values.append(values)195 196    # Concatenate all chunks197    row_indices = torch.cat(all_row_indices)198    col_indices = torch.cat(all_col_indices)199    values = torch.cat(all_values)200 201    # Create sparse tensor202    indices = torch.stack([row_indices, col_indices])203    filterq = torch.sparse_coo_tensor(204        indices, values, (3 * num_variables, num_variables)205    )206 207    # Compute filterq.T @ filterq208    return torch.sparse.mm(filterq.t(), filterq)209 210 211# Usage example:212# variable_indices = torch.tensor(...).cuda()  # Your input tensor on GPU213# result = _buildq3d_gpu(variable_indices)214 215 216def _buildq2d(variable_indices: np.ndarray):217    """218    Builds the filterq matrix for the given variables.219 220    Version for 2 dimensions.221    """222 223    num_variables = variable_indices.max() + 1224    filterq = sparse.lil_matrix((3 * num_variables, num_variables))225 226    # Pad variable_indices to simplify out-of-bounds accesses227    variable_indices = np.pad(228        variable_indices, [(0, 1), (0, 1)], mode="constant", constant_values=-1229    )230 231    coords = np.nonzero(variable_indices >= 0)232    for count, (i, j) in enumerate(zip(*coords)):233        assert variable_indices[i, j] == count234 235        filterq[2 * count, count] = -2236        neighbor = variable_indices[i - 1, j]237        if neighbor >= 0:238            filterq[2 * count, neighbor] = 1239        else:240            filterq[2 * count, count] += 1241 242        neighbor = variable_indices[i + 1, j]243        if neighbor >= 0:244            filterq[2 * count, neighbor] = 1245        else:246            filterq[2 * count, count] += 1247 248        filterq[2 * count + 1, count] = -2249        neighbor = variable_indices[i, j - 1]250        if neighbor >= 0:251            filterq[2 * count + 1, neighbor] = 1252        else:253            filterq[2 * count + 1, count] += 1254 255        neighbor = variable_indices[i, j + 1]256        if neighbor >= 0:257            filterq[2 * count + 1, neighbor] = 1258        else:259            filterq[2 * count + 1, count] += 1260 261    filterq = filterq.tocsr()262    return filterq.T.dot(filterq)263 264 265def _jacobi(266    filterq,267    x0: np.ndarray,268    lower_bound: np.ndarray,269    upper_bound: np.ndarray,270    max_iters: int = 10,271    rel_tol: float = 1e-6,272    weight: float = 0.5,273):274    """Jacobi method with constraints."""275 276    jacobi_r = sparse.lil_matrix(filterq)277    shp = jacobi_r.shape278    jacobi_d = 1.0 / filterq.diagonal()279    jacobi_r.setdiag((0,) * shp[0])280    jacobi_r = jacobi_r.tocsr()281 282    x = x0283 284    # We check the stopping criterion each 10 iterations285    check_each = 10286    cum_rel_tol = 1 - (1 - rel_tol) ** check_each287 288    energy_now = np.dot(x, filterq.dot(x)) / 2289    logging.info("Energy at iter %d: %.6g", 0, energy_now)290    for i in range(max_iters):291 292        x_1 = -jacobi_d * jacobi_r.dot(x)293        x = weight * x_1 + (1 - weight) * x294 295        # Constraints.296        x = np.maximum(x, lower_bound)297        x = np.minimum(x, upper_bound)298 299        # Stopping criterion300        if (i + 1) % check_each == 0:301            # Update energy302            energy_before = energy_now303            energy_now = np.dot(x, filterq.dot(x)) / 2304 305            logging.info("Energy at iter %d: %.6g", i + 1, energy_now)306 307            # Check stopping criterion308            cum_rel_improvement = (energy_before - energy_now) / energy_before309            if cum_rel_improvement < cum_rel_tol:310                break311 312    return x313 314 315def signed_distance_function(316    levelset: np.ndarray, band_radius: int317) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:318    """319    Return the distance to the 0.5 levelset of a function, the mask of the320    border (i.e., the nearest cells to the 0.5 level-set) and the mask of the321    band (i.e., the cells of the function whose distance to the 0.5 level-set322    is less of equal to `band_radius`).323    """324 325    binary_array = np.where(levelset > 0, True, False)326 327    # Compute the band and the border.328    dist_func = ndi.distance_transform_edt329    distance = np.where(330        binary_array, dist_func(binary_array) - 0.5, -dist_func(~binary_array) + 0.5331    )332    border = np.abs(distance) < 1333    band = np.abs(distance) <= band_radius334 335    return distance, border, band336 337 338def signed_distance_function_iso0(339    levelset: np.ndarray, band_radius: int340) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:341    """342    Return the distance to the 0 levelset of a function, the mask of the343    border (i.e., the nearest cells to the 0 level-set) and the mask of the344    band (i.e., the cells of the function whose distance to the 0 level-set345    is less of equal to `band_radius`).346    """347 348    binary_array = levelset > 0349 350    # Compute the band and the border.351    dist_func = ndi.distance_transform_edt352    distance = np.where(353        binary_array, dist_func(binary_array), -dist_func(~binary_array)354    )355    border = np.zeros_like(levelset, dtype=bool)356    border[:-1, :, :] |= levelset[:-1, :, :] * levelset[1:, :, :] <= 0357    border[:, :-1, :] |= levelset[:, :-1, :] * levelset[:, 1:, :] <= 0358    border[:, :, :-1] |= levelset[:, :, :-1] * levelset[:, :, 1:] <= 0359    band = np.abs(distance) <= band_radius360 361    return distance, border, band362 363 364def signed_distance_function_gpu(levelset: torch.Tensor, band_radius: int):365    binary_array = (levelset > 0).float()366 367    # Compute distance transform368    dist_pos = (369        F.max_pool3d(370            -binary_array.unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1371        )372        .squeeze(0)373        .squeeze(0)374        + binary_array375    )376    dist_neg = F.max_pool3d(377        (binary_array - 1).unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1378    ).squeeze(0).squeeze(0) + (1 - binary_array)379 380    distance = torch.where(binary_array > 0, dist_pos - 0.5, -dist_neg + 0.5)381 382    # breakpoint()383 384    # Use levelset as distance directly385    # distance = levelset386    # print(distance.shape)387    # Compute border and band388    border = torch.abs(distance) < 1389    band = torch.abs(distance) <= band_radius390 391    return distance, border, band392 393 394def smooth_constrained(395    binary_array: np.ndarray,396    band_radius: int = 4,397    max_iters: int = 250,398    rel_tol: float = 1e-6,399) -> np.ndarray:400    """401    Implementation of the smoothing method from402 403    "Surface Extraction from Binary Volumes with Higher-Order Smoothness"404    Victor Lempitsky, CVPR10405    """406 407    # # Compute the distance map, the border and the band.408    logging.info("Computing distance transform...")409    # distance, _, band = signed_distance_function(binary_array, band_radius)410    binary_array_gpu = torch.from_numpy(binary_array).cuda()411    distance, _, band = signed_distance_function_gpu(binary_array_gpu, band_radius)412    distance = distance.cpu().numpy()413    band = band.cpu().numpy()414 415    variable_indices = _build_variable_indices(band)416 417    # Compute filterq.418    logging.info("Building matrix filterq...")419    if binary_array.ndim == 3:420        filterq = _buildq3d(variable_indices)421        # variable_indices_gpu = torch.from_numpy(variable_indices).cuda()422        # filterq_gpu = _buildq3d_gpu(variable_indices_gpu)423        # filterq = filterq_gpu.cpu().numpy()424    elif binary_array.ndim == 2:425        filterq = _buildq2d(variable_indices)426    else:427        raise ValueError("binary_array.ndim not in [2, 3]")428 429    # Initialize the variables.430    res = np.asarray(distance, dtype=np.double)431    x = res[band]432    upper_bound = np.where(x < 0, x, np.inf)433    lower_bound = np.where(x > 0, x, -np.inf)434 435    upper_bound[np.abs(upper_bound) < 1] = 0436    lower_bound[np.abs(lower_bound) < 1] = 0437 438    # Solve.439    logging.info("Minimizing energy...")440    x = _jacobi(441        filterq=filterq,442        x0=x,443        lower_bound=lower_bound,444        upper_bound=upper_bound,445        max_iters=max_iters,446        rel_tol=rel_tol,447    )448 449    res[band] = x450    return res451 452 453def total_variation_denoising(x, weight=0.1, num_iterations=5, eps=1e-8):454    diff_x = torch.diff(x, dim=0, prepend=x[:1])455    diff_y = torch.diff(x, dim=1, prepend=x[:, :1])456    diff_z = torch.diff(x, dim=2, prepend=x[:, :, :1])457 458    norm = torch.sqrt(diff_x**2 + diff_y**2 + diff_z**2 + eps)459 460    div_x = torch.diff(diff_x / norm, dim=0, append=diff_x[-1:] / norm[-1:])461    div_y = torch.diff(diff_y / norm, dim=1, append=diff_y[:, -1:] / norm[:, -1:])462    div_z = torch.diff(diff_z / norm, dim=2, append=diff_z[:, :, -1:] / norm[:, :, -1:])463 464    return x - weight * (div_x + div_y + div_z)465 466 467def smooth_constrained_gpu(468    binary_array: torch.Tensor,469    band_radius: int = 4,470    max_iters: int = 250,471    rel_tol: float = 1e-4,472):473    distance, _, band = signed_distance_function_gpu(binary_array, band_radius)474 475    # Initialize variables476    x = distance[band]477    upper_bound = torch.where(x < 0, x, torch.tensor(float("inf"), device=x.device))478    lower_bound = torch.where(x > 0, x, torch.tensor(float("-inf"), device=x.device))479 480    upper_bound[torch.abs(upper_bound) < 1] = 0481    lower_bound[torch.abs(lower_bound) < 1] = 0482 483    # Define the 3D Laplacian kernel484    laplacian_kernel = torch.tensor(485        [486            [487                [488                    [[0, 1, 0], [1, -6, 1], [0, 1, 0]],489                    [[1, 0, 1], [0, 0, 0], [1, 0, 1]],490                    [[0, 1, 0], [1, 0, 1], [0, 1, 0]],491                ]492            ]493        ],494        device=x.device,495    ).float()496 497    laplacian_kernel = laplacian_kernel / laplacian_kernel.abs().sum()498 499    breakpoint()500 501    # Simplified Jacobi iteration502    for i in range(max_iters):503        # Reshape x to 5D tensor (batch, channel, depth, height, width)504        x_5d = x.view(1, 1, *band.shape)505        x_3d = x.view(*band.shape)506 507        # Apply 3D convolution508        laplacian = F.conv3d(x_5d, laplacian_kernel, padding=1)509 510        # Reshape back to original dimensions511        laplacian = laplacian.view(x.shape)512 513        # Use a small relaxation factor to improve stability514        relaxation_factor = 0.1515        tv_weight = 0.1516        # x_new = x + relaxation_factor * laplacian517        x_new = total_variation_denoising(x_3d, weight=tv_weight)518        # Print laplacian min and max519        # print(f"Laplacian min: {laplacian.min().item():.4f}, max: {laplacian.max().item():.4f}")520 521        # Apply constraints522        # Reshape x_new to match the dimensions of lower_bound and upper_bound523        x_new = x_new.view(x.shape)524        x_new = torch.clamp(x_new, min=lower_bound, max=upper_bound)525 526        # Check for convergence527        diff_norm = torch.norm(x_new - x)528        print(diff_norm)529        x_norm = torch.norm(x)530 531        if x_norm > 1e-8:  # Avoid division by very small numbers532            relative_change = diff_norm / x_norm533            if relative_change < rel_tol:534                break535        elif diff_norm < rel_tol:  # If x_norm is very small, check absolute change536            break537 538        x = x_new539 540        # Check for NaN and break if found, also check for inf541        if torch.isnan(x).any() or torch.isinf(x).any():542            print(f"NaN or Inf detected at iteration {i}")543            breakpoint()544            break545 546    result = distance.clone()547    result[band] = x548    return result549 550 551def smooth_gaussian(binary_array: np.ndarray, sigma: float = 3) -> np.ndarray:552    vol = np.float_(binary_array) - 0.5553    return ndi.gaussian_filter(vol, sigma=sigma)554 555 556def smooth_gaussian_gpu(binary_array: torch.Tensor, sigma: float = 3):557    # vol = binary_array.float()558    vol = binary_array559    kernel_size = int(2 * sigma + 1)560    kernel = torch.ones(561        1,562        1,563        kernel_size,564        kernel_size,565        kernel_size,566        device=binary_array.device,567        dtype=vol.dtype,568    ) / (kernel_size**3)569    return F.conv3d(570        vol.unsqueeze(0).unsqueeze(0), kernel, padding=kernel_size // 2571    ).squeeze()572 573 574def smooth(binary_array: np.ndarray, method: str = "auto", **kwargs) -> np.ndarray:575    """576    Smooths the 0.5 level-set of a binary array. Returns a floating-point577    array with a smoothed version of the original level-set in the 0 isovalue.578 579    This function can apply two different methods:580 581    - A constrained smoothing method which preserves details and fine582      structures, but it is slow and requires a large amount of memory. This583      method is recommended when the input array is small (smaller than584      (500, 500, 500)).585    - A Gaussian filter applied over the binary array. This method is fast, but586      not very precise, as it can destroy fine details. It is only recommended587      when the input array is large and the 0.5 level-set does not contain588      thin structures.589 590    Parameters591    ----------592    binary_array : ndarray593        Input binary array with the 0.5 level-set to smooth.594    method : str, one of ['auto', 'gaussian', 'constrained']595        Smoothing method. If 'auto' is given, the method will be automatically596        chosen based on the size of `binary_array`.597 598    Parameters for 'gaussian'599    -------------------------600    sigma : float601        Size of the Gaussian filter (default 3).602 603    Parameters for 'constrained'604    ----------------------------605    max_iters : positive integer606        Number of iterations of the constrained optimization method607        (default 250).608    rel_tol: float609        Relative tolerance as a stopping criterion (default 1e-6).610 611    Output612    ------613    res : ndarray614        Floating-point array with a smoothed 0 level-set.615    """616 617    binary_array = np.asarray(binary_array)618 619    if method == "auto":620        if binary_array.size > 512**3:621            method = "gaussian"622        else:623            method = "constrained"624 625    if method == "gaussian":626        return smooth_gaussian(binary_array, **kwargs)627 628    if method == "constrained":629        return smooth_constrained(binary_array, **kwargs)630 631    raise ValueError("Unknown method '{}'".format(method))632 633 634def smooth_gpu(binary_array: torch.Tensor, method: str = "auto", **kwargs):635    if method == "auto":636        method = "gaussian" if binary_array.numel() > 512**3 else "constrained"637 638    if method == "gaussian":639        return smooth_gaussian_gpu(binary_array, **kwargs)640    elif method == "constrained":641        return smooth_constrained_gpu(binary_array, **kwargs)642    else:643        raise ValueError(f"Unknown method '{method}'")644