init

lidm/eval/README.md

@@ -1,95 +0,0 @@
-# Evaluation Toolbox for LiDAR Generation
-
-This directory is a **self-contained**, **memory-friendly** and mostly **CUDA-accelerated** toolbox of multiple evaluation metrics for LiDAR generative models, including:
-* Perceptual metrics (our proposed):
-  * Fréchet Range Image Distance (**FRID**)
-  * Fréchet Sparse Volume Distance (**FSVD**)
-  * Fréchet Point-based Volume Distance (**FPVD**)
-* Statistical metrics (proposed in [Learning Representations and Generative Models for 3D Point Clouds](https://arxiv.org/abs/1707.02392)):
-  * Minimum Matching Distance (**MMD**)
-  * Jensen-Shannon Divergence (**JSD**)
-* Statistical pairwise metrics (for reconstruction only):
-  * Chamfer Distance (**CD**)
-  * Earth Mover's Distance (**EMD**)
-
-## Citation
-
-If you find this project useful in your research, please consider citing:
-```
-@article{ran2024towards,
-  title={Towards Realistic Scene Generation with LiDAR Diffusion Models},
-  author={Ran, Haoxi and Guizilini, Vitor and Wang, Yue},
-  journal={arXiv preprint arXiv:2404.00815},
-  year={2024}
-}
-```
-
-
-## Dependencies
-
-### Basic (install through **pip**):
-* scipy
-* numpy
-* torch
-* pyyaml
-
-### Required by FSVD and FPVD:
-* [Torchsparse v1.4.0](https://github.com/mit-han-lab/torchsparse/tree/v1.4.0) (pip install git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0)
-* [Google Sparse Hash library](https://github.com/sparsehash/sparsehash) (apt-get install libsparsehash-dev **or** compile locally and update variable CPLUS_INCLUDE_PATH with directory path)
-
-
-## Model Zoo 
-
-To evaluate with perceptual metrics on different types of LiDAR data, you can download all models through:
-*  this [google drive link](https://drive.google.com/file/d/1Ml4p4_nMlwLkSp7JB528GJv2_HxO8v1i/view?usp=drive_link) in the .zip file 
-
-or
-*  the **full directory** of one specific model:
-
-### 64-beam LiDAR (trained on [SemanticKITTI](http://semantic-kitti.org/dataset.html)):
-
-| Metric |                                            Model                                            |          Arch           |                                                  Link                                                   | Code                                                             | Comments                                                                  |
-|:------:|:-------------------------------------------------------------------------------------------:|:-----------------------:|:-------------------------------------------------------------------------------------------------------:|:-----------------------------------------------------------------|---------------------------------------------------------------------------|
-|  FRID  | [RangeNet++](https://www.ipb.uni-bonn.de/wp-content/papercite-data/pdf/milioto2019iros.pdf) |  DarkNet21-based UNet   | [Google Drive](https://drive.google.com/drive/folders/1ZS8KOoxB9hjB6kwKbH5Zfc8O5qJlKsbl?usp=drive_link) | [./models/rangenet/model.py](./models/rangenet/model.py)         | range image input (our trained model without the need of remission input) |
-|  FSVD  |                      [MinkowskiNet](https://arxiv.org/abs/1904.08755)                       |       Sparse UNet       | [Google Drive](https://drive.google.com/drive/folders/1zN12ZEvjIvo4PCjAsncgC22yvtRrCCMe?usp=drive_link) | [./models/minkowskinet/model.py](./models/minkowskinet/model.py) | point cloud input                                                         |
-|  FPVD  |                         [SPVCNN](https://arxiv.org/abs/2007.16100)                          | Point-Voxel Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1oEm3qpxfGetiVAfXIvecawEiFqW79M6B?usp=drive_link) | [./models/spvcnn/model.py](./models/spvcnn/model.py)             | point cloud input                                                         |
-
-
-### 32-beam LiDAR (trained on [nuScenes](https://www.nuscenes.org/nuscenes)):
-
-| Metric |                      Model                       |          Arch           |                                                  Link                                                   | Code                                                             | Comments          |
-|:------:|:------------------------------------------------:|:-----------------------:|:-------------------------------------------------------------------------------------------------------:|:-----------------------------------------------------------------|-------------------|
-|  FSVD  | [MinkowskiNet](https://arxiv.org/abs/1904.08755) |       Sparse UNet       | [Google Drive](https://drive.google.com/drive/folders/1oZIS9FlklCQ6dlh3TZ8Junir7QwgT-Me?usp=drive_link) | [./models/minkowskinet/model.py](./models/minkowskinet/model.py) | point cloud input |
-|  FPVD  |    [SPVCNN](https://arxiv.org/abs/2007.16100)    | Point-Voxel Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1F69RbprAoT6MOJ7iI0KHjxuq-tbeqGiR?usp=drive_link) | [./models/spvcnn/model.py](./models/spvcnn/model.py)             | point cloud input |
-
-
-## Usage
-
-1. Place the unzipped `pretrained_weights` folder under the root python directory **or** modify the `DEFAULT_ROOT` variable in the `__init__.py`.
-2. Prepare input data, including the synthesized samples and the reference dataset. **Note**: The reference data should be the **point clouds projected back from range images** instead of raw point clouds. 
-3. Specify the data type (`32` or `64`) and the metrics to evaluate. Options: `mmd`, `jsd`, `frid`, `fsvd`, `fpvd`, `cd`, `emd`.
-4. (Optional) If you want to compute `frid`, `fsvd` or `fpvd` metric, adjust the corresponding batch size through the `MODAL2BATCHSIZE` in file `__init__.py` according to your max GPU memory (default: ~24GB).
-5. Start evaluation and all results will print out!
-
-### Example:
-
-```
-from .eval_utils import evaluate
-
-data = '64'  # specify data type to evaluate
-metrics = ['mmd', 'jsd', 'frid', 'fsvd', 'fpvd']  # specify metrics to evaluate
-
-# list of np.float32 array
-# shape of each array: (#points, #dim=3), #dim: xyz coordinate (NOTE: no need to input remission)
-reference = ...
-samples = ...
-
-evaluate(reference, samples, metrics, data)
-```
-
-
-## Acknowledgement
-
-- The implementation of MinkowskiNet and SPVCNN is borrowed from [2DPASS](https://github.com/yanx27/2DPASS).
-- The implementation of RangeNet++ is borrowed from [the official RangeNet++ codebase](https://github.com/PRBonn/lidar-bonnetal).
-- The implementation of Chamfer Distance is adapted from [CD Pytorch Implementation](https://github.com/ThibaultGROUEIX/ChamferDistancePytorch) and Earth Mover's Distance from [MSN official repo](https://github.com/Colin97/MSN-Point-Cloud-Completion).

lidm/eval/__init__.py

@@ -1,62 +0,0 @@
-"""
-@Author: Haoxi Ran
-@Date: 01/03/2024
-@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
-
-"""
-
-import os
-
-import torch
-import yaml
-
-from lidm.utils.misc_utils import dict2namespace
-from ..modules.rangenet.model import Model as rangenet
-
-try:
-    from ..modules.spvcnn.model import Model as spvcnn
-    from ..modules.minkowskinet.model import Model as minkowskinet
-except:
-    print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
-
-# user settings
-DEFAULT_ROOT = './pretrained_weights'
-MODAL2BATCHSIZE = {'range': 100, 'voxel': 50, 'point_voxel': 25}
-OUTPUT_TEMPLATE = 50 * '-' + '\n|' + 16 * ' ' + '{}:{:.4E}' + 17 * ' ' + '|\n' + 50 * '-'
-
-# eval settings (do not modify)
-VOXEL_SIZE = 0.05
-NUM_SECTORS = 16
-AGG_TYPE = 'depth'
-TYPE2DATASET = {'32': 'nuscenes', '64': 'kitti'}
-DATA_CONFIG = {'64': {'x': [-50, 50], 'y': [-50, 50], 'z': [-3, 1]},
-               '32': {'x': [-30, 30], 'y': [-30, 30], 'z': [-3, 6]}}
-MODALITY2MODEL = {'range': 'rangenet', 'voxel': 'minkowskinet', 'point_voxel': 'spvcnn'}
-DATASET_CONFIG = {'kitti': {'size': [64, 1024], 'fov': [3, -25], 'depth_range': [1.0, 56.0], 'depth_scale': 6},
-                  'nuscenes': {'size': [32, 1024], 'fov': [10, -30], 'depth_range': [1.0, 45.0]}}
-
-
-def build_model(dataset_name, model_name, device='cpu'):
-    # config
-    model_folder = os.path.join(DEFAULT_ROOT, dataset_name, model_name)
-
-    if not os.path.isdir(model_folder):
-        raise Exception('Not Available Pretrained Weights!')
-
-    config = yaml.safe_load(open(os.path.join(model_folder, 'config.yaml'), 'r'))
-    if model_name != 'rangenet':
-        config = dict2namespace(config)
-
-    # build model
-    model = eval(model_name)(config)
-
-    # load checkpoint
-    if model_name == 'rangenet':
-        model.load_pretrained_weights(model_folder)
-    else:
-        ckpt = torch.load(os.path.join(model_folder, 'model.ckpt'), map_location="cpu")
-        model.load_state_dict(ckpt['state_dict'], strict=False)
-    model.to(device)
-    model.eval()
-
-    return model

lidm/eval/compile.sh

@@ -1,9 +0,0 @@
-#!/bin/sh
-
-cd modules/chamfer
-python setup.py build_ext --inplace
-
-cd ../emd
-python setup.py build_ext --inplace
-
-cd ..

lidm/eval/eval_utils.py

@@ -1,138 +0,0 @@
-"""
-@Author: Haoxi Ran
-@Date: 01/03/2024
-@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
-
-"""
-import multiprocessing
-from functools import partial
-
-import numpy as np
-from scipy.spatial.distance import jensenshannon
-from tqdm import tqdm
-
-from . import OUTPUT_TEMPLATE
-from .metric_utils import compute_logits, compute_pairwise_cd, \
-    compute_pairwise_emd, pcd2bev_sum, compute_pairwise_cd_batch, pcd2bev_bin
-from .fid_score import calculate_frechet_distance
-
-
-def evaluate(reference, samples, metrics, data):
-    # perceptual
-    if 'frid' in metrics:
-        compute_frid(reference, samples, data)
-    if 'fsvd' in metrics:
-        compute_fsvd(reference, samples, data)
-    if 'fpvd' in metrics:
-        compute_fpvd(reference, samples, data)
-
-    # reconstruction
-    if 'cd' in metrics:
-        compute_cd(reference, samples)
-    if 'emd' in metrics:
-        compute_emd(reference, samples)
-
-    # statistical
-    if 'jsd' in metrics:
-        compute_jsd(reference, samples, data)
-    if 'mmd' in metrics:
-        compute_mmd(reference, samples, data)
-
-
-def compute_cd(reference, samples):
-    """
-    Calculate score of Chamfer Distance (CD)
-
-    """
-    print('Evaluating (CD) ...')
-    results = []
-    for x, y in zip(reference, samples):
-        d = compute_pairwise_cd(x, y)
-        results.append(d)
-    score = sum(results) / len(results)
-    print(OUTPUT_TEMPLATE.format('CD  ', score))
-
-
-def compute_emd(reference, samples):
-    """
-    Calculate score of Earth Mover's Distance (EMD)
-
-    """
-    print('Evaluating (EMD) ...')
-    results = []
-    for x, y in zip(reference, samples):
-        d = compute_pairwise_emd(x, y)
-        results.append(d)
-    score = sum(results) / len(results)
-    print(OUTPUT_TEMPLATE.format('EMD ', score))
-
-
-def compute_mmd(reference, samples, data, dist='cd', verbose=True):
-    """
-    Calculate the score of Minimum Matching Distance (MMD)
-
-    """
-    print('Evaluating (MMD) ...')
-    assert dist in ['cd', 'emd']
-    reference, samples = pcd2bev_bin(data, reference, samples)
-    compute_dist_func = compute_pairwise_cd_batch if dist == 'cd' else compute_pairwise_emd
-    results = []
-    for r in tqdm(reference, disable=not verbose):
-        dists = compute_dist_func(r, samples)
-        results.append(min(dists))
-    score = sum(results) / len(results)
-    print(OUTPUT_TEMPLATE.format('MMD ', score))
-
-
-def compute_jsd(reference, samples, data):
-    """
-    Calculate the score of Jensen-Shannon Divergence (JSD)
-
-    """
-    print('Evaluating (JSD) ...')
-    reference, samples = pcd2bev_sum(data, reference, samples)
-    reference = (reference / np.sum(reference)).flatten()
-    samples = (samples / np.sum(samples)).flatten()
-    score = jensenshannon(reference, samples)
-    print(OUTPUT_TEMPLATE.format('JSD ', score))
-
-
-def compute_fd(reference, samples):
-    mu1, mu2 = np.mean(reference, axis=0), np.mean(samples, axis=0)
-    sigma1, sigma2 = np.cov(reference, rowvar=False), np.cov(samples, rowvar=False)
-    distance = calculate_frechet_distance(mu1, sigma1, mu2, sigma2)
-    return distance
-
-
-def compute_frid(reference, samples, data):
-    """
-    Calculate the score of Fréchet Range Image Distance (FRID)
-
-    """
-    print('Evaluating (FRID) ...')
-    gt_logits, samples_logits = compute_logits(data, 'range', reference, samples)
-    score = compute_fd(gt_logits, samples_logits)
-    print(OUTPUT_TEMPLATE.format('FRID', score))
-
-
-def compute_fsvd(reference, samples, data):
-    """
-    Calculate the score of Fréchet Sparse Volume Distance (FSVD)
-
-    """
-    print('Evaluating (FSVD) ...')
-    gt_logits, samples_logits = compute_logits(data, 'voxel', reference, samples)
-    score = compute_fd(gt_logits, samples_logits)
-    print(OUTPUT_TEMPLATE.format('FSVD', score))
-
-
-def compute_fpvd(reference, samples, data):
-    """
-    Calculate the score of Fréchet Point-based Volume Distance (FPVD)
-
-    """
-    print('Evaluating (FPVD) ...')
-    gt_logits, samples_logits = compute_logits(data, 'point_voxel', reference, samples)
-    score = compute_fd(gt_logits, samples_logits)
-    print(OUTPUT_TEMPLATE.format('FPVD', score))
-

lidm/eval/fid_score.py

@@ -1,191 +0,0 @@
-"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
-The FID metric calculates the distance between two distributions of images.
-Typically, we have summary statistics (mean & covariance matrix) of one
-of these distributions, while the 2nd distribution is given by a GAN.
-When run as a stand-alone program, it compares the distribution of
-images that are stored as PNG/JPEG at a specified location with a
-distribution given by summary statistics (in pickle format).
-The FID is calculated by assuming that X_1 and X_2 are the activations of
-the pool_3 layer of the inception net for generated samples and real world
-samples respectively.
-See --help to see further details.
-Code adapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
-of Tensorflow
-Copyright 2018 Institute of Bioinformatics, JKU Linz
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-   http://www.apache.org/licenses/LICENSE-2.0
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-"""
-import os
-import pathlib
-from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
-
-import numpy as np
-import torch
-import torchvision.transforms as TF
-from PIL import Image
-from scipy import linalg
-from torch.nn.functional import adaptive_avg_pool2d
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    # If tqdm is not available, provide a mock version of it
-    def tqdm(x):
-        return x
-
-class ImagePathDataset(torch.utils.data.Dataset):
-    def __init__(self, files, transforms=None):
-        self.files = files
-        self.transforms = transforms
-
-    def __len__(self):
-        return len(self.files)
-
-    def __getitem__(self, i):
-        path = self.files[i]
-        img = Image.open(path).convert('RGB')
-        if self.transforms is not None:
-            img = self.transforms(img)
-        return img
-
-
-def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
-                    num_workers=1):
-    """Calculates the activations of the pool_3 layer for all images.
-    Params:
-    -- files       : List of image files paths
-    -- model       : Instance of inception model
-    -- batch_size  : Batch size of images for the model to process at once.
-                     Make sure that the number of samples is a multiple of
-                     the batch size, otherwise some samples are ignored. This
-                     behavior is retained to match the original FID score
-                     implementation.
-    -- dims        : Dimensionality of features returned by Inception
-    -- device      : Device to run calculations
-    -- num_workers : Number of parallel dataloader workers
-    Returns:
-    -- A numpy array of dimension (num images, dims) that contains the
-       activations of the given tensor when feeding inception with the
-       query tensor.
-    """
-    model.eval()
-
-    if batch_size > len(files):
-        print(('Warning: batch size is bigger than the data size. '
-               'Setting batch size to data size'))
-        batch_size = len(files)
-
-    dataset = ImagePathDataset(files, transforms=TF.ToTensor())
-    dataloader = torch.utils.data.DataLoader(dataset,
-                                             batch_size=batch_size,
-                                             shuffle=False,
-                                             drop_last=False,
-                                             num_workers=num_workers)
-
-    pred_arr = np.empty((len(files), dims))
-
-    start_idx = 0
-
-    for batch in tqdm(dataloader):
-        batch = batch.to(device)
-
-        with torch.no_grad():
-            pred = model(batch)[0]
-
-        # If model output is not scalar, apply global spatial average pooling.
-        # This happens if you choose a dimensionality not equal 2048.
-        if pred.size(2) != 1 or pred.size(3) != 1:
-            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
-
-        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
-
-        pred_arr[start_idx:start_idx + pred.shape[0]] = pred
-
-        start_idx = start_idx + pred.shape[0]
-
-    return pred_arr
-
-
-def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
-    """Numpy implementation of the Frechet Distance.
-    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
-    and X_2 ~ N(mu_2, C_2) is
-            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
-    Stable version by Dougal J. Sutherland.
-    Params:
-    -- mu1   : Numpy array containing the activations of a layer of the
-               inception net (like returned by the function 'get_predictions')
-               for generated samples.
-    -- mu2   : The sample mean over activations, precalculated on an
-               representative data set.
-    -- sigma1: The covariance matrix over activations for generated samples.
-    -- sigma2: The covariance matrix over activations, precalculated on an
-               representative data set.
-    Returns:
-    --   : The Frechet Distance.
-    """
-
-    mu1 = np.atleast_1d(mu1)
-    mu2 = np.atleast_1d(mu2)
-
-    sigma1 = np.atleast_2d(sigma1)
-    sigma2 = np.atleast_2d(sigma2)
-
-    assert mu1.shape == mu2.shape, \
-        'Training and test mean vectors have different lengths'
-    assert sigma1.shape == sigma2.shape, \
-        'Training and test covariances have different dimensions'
-
-    diff = mu1 - mu2
-
-    # Product might be almost singular
-    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
-    if not np.isfinite(covmean).all():
-        msg = ('fid calculation produces singular product; '
-               'adding %s to diagonal of cov estimates') % eps
-        print(msg)
-        offset = np.eye(sigma1.shape[0]) * eps
-        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
-
-    # Numerical error might give slight imaginary component
-    if np.iscomplexobj(covmean):
-        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
-            m = np.max(np.abs(covmean.imag))
-            raise ValueError('Imaginary component {}'.format(m))
-        covmean = covmean.real
-
-    tr_covmean = np.trace(covmean)
-
-    return (diff.dot(diff) + np.trace(sigma1)
-            + np.trace(sigma2) - 2 * tr_covmean)
-
-
-def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
-                                    device='cpu', num_workers=1):
-    """Calculation of the statistics used by the FID.
-    Params:
-    -- files       : List of image files paths
-    -- model       : Instance of inception model
-    -- batch_size  : The images numpy array is split into batches with
-                     batch size batch_size. A reasonable batch size
-                     depends on the hardware.
-    -- dims        : Dimensionality of features returned by Inception
-    -- device      : Device to run calculations
-    -- num_workers : Number of parallel dataloader workers
-    Returns:
-    -- mu    : The mean over samples of the activations of the pool_3 layer of
-               the inception model.
-    -- sigma : The covariance matrix of the activations of the pool_3 layer of
-               the inception model.
-    """
-    act = get_activations(files, model, batch_size, dims, device, num_workers)
-    mu = np.mean(act, axis=0)
-    sigma = np.cov(act, rowvar=False)
-    return mu, sigma

lidm/eval/metric_utils.py

@@ -1,458 +0,0 @@
-"""
-@Author: Haoxi Ran
-@Date: 01/03/2024
-@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
-
-"""
-
-import math
-from itertools import repeat
-from typing import List, Tuple, Union
-import numpy as np
-import torch
-
-from . import build_model, VOXEL_SIZE, MODALITY2MODEL, MODAL2BATCHSIZE, DATASET_CONFIG, AGG_TYPE, NUM_SECTORS, \
-    TYPE2DATASET, DATA_CONFIG
-
-try:
-    from torchsparse import SparseTensor, PointTensor
-    from torchsparse.utils.collate import sparse_collate_fn
-    from .modules.chamfer3D.dist_chamfer_3D import chamfer_3DDist
-    from .modules.chamfer2D.dist_chamfer_2D import chamfer_2DDist
-    from .modules.emd.emd_module import emdModule
-except:
-    print(
-        'To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
-
-
-def ravel_hash(x: np.ndarray) -> np.ndarray:
-    assert x.ndim == 2, x.shape
-
-    x = x - np.min(x, axis=0)
-    x = x.astype(np.uint64, copy=False)
-    xmax = np.max(x, axis=0).astype(np.uint64) + 1
-
-    h = np.zeros(x.shape[0], dtype=np.uint64)
-    for k in range(x.shape[1] - 1):
-        h += x[:, k]
-        h *= xmax[k + 1]
-    h += x[:, -1]
-    return h
-
-
-def sparse_quantize(coords, voxel_size: Union[float, Tuple[float, ...]] = 1, *, return_index: bool = False,
-                    return_inverse: bool = False) -> List[np.ndarray]:
-    """
-    Modified based on https://github.com/mit-han-lab/torchsparse/blob/462dea4a701f87a7545afb3616bf2cf53dd404f3/torchsparse/utils/quantize.py
-
-    """
-    if isinstance(voxel_size, (float, int)):
-        voxel_size = tuple(repeat(voxel_size, coords.shape[1]))
-    assert isinstance(voxel_size, tuple) and len(voxel_size) in [2, 3]  # support 2D and 3D coordinates only
-
-    voxel_size = np.array(voxel_size)
-    coords = np.floor(coords / voxel_size).astype(np.int32)
-
-    _, indices, inverse_indices = np.unique(
-        ravel_hash(coords), return_index=True, return_inverse=True
-    )
-    coords = coords[indices]
-
-    outputs = [coords]
-    if return_index:
-        outputs += [indices]
-    if return_inverse:
-        outputs += [inverse_indices]
-    return outputs[0] if len(outputs) == 1 else outputs
-
-
-def pcd2range(pcd, size, fov, depth_range, remission=None, labels=None, **kwargs):
-    # laser parameters
-    fov_up = fov[0] / 180.0 * np.pi  # field of view up in rad
-    fov_down = fov[1] / 180.0 * np.pi  # field of view down in rad
-    fov_range = abs(fov_down) + abs(fov_up)  # get field of view total in rad
-
-    # get depth (distance) of all points
-    depth = np.linalg.norm(pcd, 2, axis=1)
-
-    # mask points out of range
-    mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
-    depth, pcd = depth[mask], pcd[mask]
-
-    # get scan components
-    scan_x, scan_y, scan_z = pcd[:, 0], pcd[:, 1], pcd[:, 2]
-
-    # get angles of all points
-    yaw = -np.arctan2(scan_y, scan_x)
-    pitch = np.arcsin(scan_z / depth)
-
-    # get projections in image coords
-    proj_x = 0.5 * (yaw / np.pi + 1.0)  # in [0.0, 1.0]
-    proj_y = 1.0 - (pitch + abs(fov_down)) / fov_range  # in [0.0, 1.0]
-
-    # scale to image size using angular resolution
-    proj_x *= size[1]  # in [0.0, W]
-    proj_y *= size[0]  # in [0.0, H]
-
-    # round and clamp for use as index
-    proj_x = np.maximum(0, np.minimum(size[1] - 1, np.floor(proj_x))).astype(np.int32)  # in [0,W-1]
-    proj_y = np.maximum(0, np.minimum(size[0] - 1, np.floor(proj_y))).astype(np.int32)  # in [0,H-1]
-
-    # order in decreasing depth
-    order = np.argsort(depth)[::-1]
-    proj_x, proj_y = proj_x[order], proj_y[order]
-
-    # project depth
-    depth = depth[order]
-    proj_range = np.full(size, -1, dtype=np.float32)
-    proj_range[proj_y, proj_x] = depth
-
-    # project point feature
-    if remission is not None:
-        remission = remission[mask][order]
-        proj_feature = np.full(size, -1, dtype=np.float32)
-        proj_feature[proj_y, proj_x] = remission
-    elif labels is not None:
-        labels = labels[mask][order]
-        proj_feature = np.full(size, 0, dtype=np.float32)
-        proj_feature[proj_y, proj_x] = labels
-    else:
-        proj_feature = None
-
-    return proj_range, proj_feature
-
-
-def range2xyz(range_img, fov, depth_range, depth_scale, log_scale=True, **kwargs):
-    # laser parameters
-    size = range_img.shape
-    fov_up = fov[0] / 180.0 * np.pi  # field of view up in rad
-    fov_down = fov[1] / 180.0 * np.pi  # field of view down in rad
-    fov_range = abs(fov_down) + abs(fov_up)  # get field of view total in rad
-
-    # inverse transform from depth
-    if log_scale:
-        depth = (np.exp2(range_img * depth_scale) - 1)
-    else:
-        depth = range_img
-
-    scan_x, scan_y = np.meshgrid(np.arange(size[1]), np.arange(size[0]))
-    scan_x = scan_x.astype(np.float64) / size[1]
-    scan_y = scan_y.astype(np.float64) / size[0]
-
-    yaw = np.pi * (scan_x * 2 - 1)
-    pitch = (1.0 - scan_y) * fov_range - abs(fov_down)
-
-    xyz = -np.ones((3, *size))
-    xyz[0] = np.cos(yaw) * np.cos(pitch) * depth
-    xyz[1] = -np.sin(yaw) * np.cos(pitch) * depth
-    xyz[2] = np.sin(pitch) * depth
-
-    # mask out invalid points
-    mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
-    xyz[:, ~mask] = -1
-
-    return xyz
-
-
-def pcd2voxel(pcd):
-    pcd_voxel = np.round(pcd / VOXEL_SIZE)
-    pcd_voxel = pcd_voxel - pcd_voxel.min(0, keepdims=1)
-    feat = np.concatenate((pcd, -np.ones((pcd.shape[0], 1))), axis=1)  # -1 for remission placeholder
-    _, inds, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-
-    feat = torch.FloatTensor(feat[inds])
-    pcd_voxel = torch.LongTensor(pcd_voxel[inds])
-    lidar = SparseTensor(feat, pcd_voxel)
-    output = {'lidar': lidar}
-    return output
-
-
-def pcd2voxel_full(data_type, *args):
-    config = DATA_CONFIG[data_type]
-    x_range, y_range, z_range = config['x'], config['y'], config['z']
-    vol_shape = (math.ceil((x_range[1] - x_range[0]) / VOXEL_SIZE), math.ceil((y_range[1] - y_range[0]) / VOXEL_SIZE),
-                 math.ceil((z_range[1] - z_range[0]) / VOXEL_SIZE))
-    min_bound = (math.ceil((x_range[0]) / VOXEL_SIZE), math.ceil((y_range[0]) / VOXEL_SIZE),
-                 math.ceil((z_range[0]) / VOXEL_SIZE))
-
-    output = tuple()
-    for data in args:
-        volume_list = []
-        for pcd in data:
-            # mask out invalid points
-            mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
-            mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
-            mask_z = np.logical_and(pcd[:, 2] > z_range[0], pcd[:, 2] < z_range[1])
-            mask = mask_x & mask_y & mask_z
-            pcd = pcd[mask]
-
-            # voxelize
-            pcd_voxel = np.floor(pcd / VOXEL_SIZE)
-            _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-            pcd_voxel = pcd_voxel[indices]
-            pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
-
-            # 2D bev grid
-            vol = np.zeros(vol_shape, dtype=np.float32)
-            vol[pcd_voxel[:, 0], pcd_voxel[:, 1], pcd_voxel[:, 2]] = 1
-            volume_list.append(vol)
-        output += (volume_list,)
-    return output
-
-
-# def pcd2bev_full(data_type, *args, voxel_size=VOXEL_SIZE):
-#     config = DATA_CONFIG[data_type]
-#     x_range, y_range = config['x'], config['y']
-#     vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
-#     min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
-#
-#     output = tuple()
-#     for data in args:
-#         volume_list = []
-#         for pcd in data:
-#             # mask out invalid points
-#             mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
-#             mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
-#             mask = mask_x & mask_y
-#             pcd = pcd[mask][:, :2]  # keep x,y coord
-#
-#             # voxelize
-#             pcd_voxel = np.floor(pcd / voxel_size)
-#             _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-#             pcd_voxel = pcd_voxel[indices]
-#             pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
-#
-#             # 2D bev grid
-#             vol = np.zeros(vol_shape, dtype=np.float32)
-#             vol[pcd_voxel[:, 0], pcd_voxel[:, 1]] = 1
-#             volume_list.append(vol)
-#         output += (volume_list,)
-#     return output
-
-
-def pcd2bev_sum(data_type, *args, voxel_size=VOXEL_SIZE):
-    config = DATA_CONFIG[data_type]
-    x_range, y_range = config['x'], config['y']
-    vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
-    min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
-
-    output = tuple()
-    for data in args:
-        volume_sum = np.zeros(vol_shape, np.float32)
-        for pcd in data:
-            # mask out invalid points
-            mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
-            mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
-            mask = mask_x & mask_y
-            pcd = pcd[mask][:, :2]  # keep x,y coord
-
-            # voxelize
-            pcd_voxel = np.floor(pcd / voxel_size)
-            _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-            pcd_voxel = pcd_voxel[indices]
-            pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
-
-            # summation
-            volume_sum[pcd_voxel[:, 0], pcd_voxel[:, 1]] += 1.
-        output += (volume_sum,)
-    return output
-
-
-def pcd2bev_bin(data_type, *args, voxel_size=0.5):
-    config = DATA_CONFIG[data_type]
-    x_range, y_range = config['x'], config['y']
-    vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
-    min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
-
-    output = tuple()
-    for data in args:
-        pcd_list = []
-        for pcd in data:
-            # mask out invalid points
-            mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
-            mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
-            mask = mask_x & mask_y
-            pcd = pcd[mask][:, :2]  # keep x,y coord
-
-            # voxelize
-            pcd_voxel = np.floor(pcd / voxel_size)
-            _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-            pcd_voxel = pcd_voxel[indices]
-            pcd_voxel = ((pcd_voxel - min_bound) / vol_shape).astype(np.float32)
-            pcd_list.append(pcd_voxel)
-        output += (pcd_list,)
-    return output
-
-
-def bev_sample(data_type, *args, voxel_size=0.5):
-    config = DATA_CONFIG[data_type]
-    x_range, y_range = config['x'], config['y']
-
-    output = tuple()
-    for data in args:
-        pcd_list = []
-        for pcd in data:
-            # mask out invalid points
-            mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
-            mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
-            mask = mask_x & mask_y
-            pcd = pcd[mask][:, :2]  # keep x,y coord
-
-            # voxelize
-            pcd_voxel = np.floor(pcd / voxel_size)
-            _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
-            pcd = pcd[indices]
-            pcd_list.append(pcd)
-        output += (pcd_list,)
-    return output
-
-
-def preprocess_pcd(pcd, **kwargs):
-    depth = np.linalg.norm(pcd, 2, axis=1)
-    mask = np.logical_and(depth > kwargs['depth_range'][0], depth < kwargs['depth_range'][1])
-    pcd = pcd[mask]
-    return pcd
-
-
-def preprocess_range(pcd, **kwargs):
-    depth_img = pcd2range(pcd, **kwargs)[0]
-    xyz_img = range2xyz(depth_img, log_scale=False, **kwargs)
-    depth_img = depth_img[None]
-    img = np.vstack([depth_img, xyz_img])
-    return img
-
-
-def batch2list(batch_dict, agg_type='depth', **kwargs):
-    """
-    Aggregation Type: Default 'depth', ['all', 'sector', 'depth']
-    """
-    output_list = []
-    batch_indices = batch_dict['batch_indices']
-    for b_idx in range(batch_indices.max() + 1):
-        # avg all
-        if agg_type == 'all':
-            logits = batch_dict['logits'][batch_indices == b_idx].mean(0)
-
-        # avg on sectors
-        elif agg_type == 'sector':
-            logits = batch_dict['logits'][batch_indices == b_idx]
-            coords = batch_dict['coords'][batch_indices == b_idx].float()
-            coords = coords - coords.mean(0)
-            angle = torch.atan2(coords[:, 1], coords[:, 0])  # [-pi, pi]
-            sector_range = torch.linspace(-np.pi - 1e-4, np.pi + 1e-4, NUM_SECTORS + 1)
-            logits_list = []
-            for i in range(NUM_SECTORS):
-                sector_indices = torch.where((angle >= sector_range[i]) & (angle < sector_range[i + 1]))[0]
-                sector_logits = logits[sector_indices].mean(0)
-                sector_logits = torch.nan_to_num(sector_logits, 0.)
-                logits_list.append(sector_logits)
-            logits = torch.cat(logits_list)  # dim: 768
-
-        # avg by depth
-        elif agg_type == 'depth':
-            logits = batch_dict['logits'][batch_indices == b_idx]
-            coords = batch_dict['coords'][batch_indices == b_idx].float()
-            coords = coords - coords.mean(0)
-            bev_depth = torch.norm(coords, dim=-1) * VOXEL_SIZE
-            sector_range = torch.linspace(kwargs['depth_range'][0] + 3, kwargs['depth_range'][1], NUM_SECTORS + 1)
-            sector_range[0] = 0.
-            logits_list = []
-            for i in range(NUM_SECTORS):
-                sector_indices = torch.where((bev_depth >= sector_range[i]) & (bev_depth < sector_range[i + 1]))[0]
-                sector_logits = logits[sector_indices].mean(0)
-                sector_logits = torch.nan_to_num(sector_logits, 0.)
-                logits_list.append(sector_logits)
-            logits = torch.cat(logits_list)  # dim: 768
-
-        else:
-            raise NotImplementedError
-
-        output_list.append(logits.detach().cpu().numpy())
-    return output_list
-
-
-def compute_logits(data_type, modality, *args):
-    assert data_type in ['32', '64']
-    assert modality in ['range', 'voxel', 'point_voxel']
-    is_voxel = 'voxel' in modality
-    dataset_name = TYPE2DATASET[data_type]
-    dataset_config = DATASET_CONFIG[dataset_name]
-    bs = MODAL2BATCHSIZE[modality]
-
-    model = build_model(dataset_name, MODALITY2MODEL[modality], device='cuda')
-
-    output = tuple()
-    for data in args:
-        all_logits_list = []
-        for i in range(math.ceil(len(data) / bs)):
-            batch = data[i * bs:(i + 1) * bs]
-            if is_voxel:
-                batch = [pcd2voxel(preprocess_pcd(pcd, **dataset_config)) for pcd in batch]
-                batch = sparse_collate_fn(batch)
-                batch = {k: v.cuda() if isinstance(v, (torch.Tensor, SparseTensor, PointTensor)) else v for k, v in
-                         batch.items()}
-                with torch.no_grad():
-                    batch_out = model(batch, return_final_logits=True)
-                    batch_out = batch2list(batch_out, AGG_TYPE, **dataset_config)
-                    all_logits_list.extend(batch_out)
-            else:
-                batch = [preprocess_range(pcd, **dataset_config) for pcd in batch]
-                batch = torch.from_numpy(np.stack(batch)).float().cuda()
-                with torch.no_grad():
-                    batch_out = model(batch, return_final_logits=True, agg_type=AGG_TYPE)
-                    all_logits_list.append(batch_out)
-        if is_voxel:
-            all_logits = np.stack(all_logits_list)
-        else:
-            all_logits = np.vstack(all_logits_list)
-        output += (all_logits,)
-
-    del model, batch, batch_out
-    torch.cuda.empty_cache()
-    return output
-
-
-def compute_pairwise_cd(x, y, module=None):
-    if module is None:
-        module = chamfer_3DDist()
-    if x.ndim == 2 and y.ndim == 2:
-        x, y = x[None], y[None]
-    x, y = torch.from_numpy(x).cuda(), torch.from_numpy(y).cuda()
-    dist1, dist2, _, _ = module(x, y)
-    dist = (dist1.mean() + dist2.mean()) / 2
-    return dist.item()
-
-
-def compute_pairwise_cd_batch(reference, samples):
-    ndim = reference.ndim
-    assert ndim in [2, 3]
-    module = chamfer_3DDist() if ndim == 3 else chamfer_2DDist()
-    len_r, len_s = reference.shape[0], [s.shape[0] for s in samples]
-    max_len = max([len_r] + len_s)
-    reference = torch.from_numpy(
-        np.vstack([reference, np.ones((max_len - reference.shape[0], ndim), dtype=np.float32) * 1e6])).cuda()
-    samples = [np.vstack([s, np.ones((max_len - s.shape[0], ndim), dtype=np.float32) * 1e6]) for s in samples]
-    samples = torch.from_numpy(np.stack(samples)).cuda()
-    reference = reference.expand_as(samples)
-    dist_r, dist_s, _, _ = module(reference, samples)
-
-    results = []
-    for i in range(samples.shape[0]):
-        dist1, dist2, len1, len2 = dist_r[i], dist_s[i], len_r, len_s[i]
-        dist = (dist1[:len1].mean() + dist2[:len2].mean()) / 2.
-        results.append(dist.item())
-    return results
-
-
-def compute_pairwise_emd(x, y, module=None):
-    if module is None:
-        module = emdModule()
-    n_points = min(x.shape[0], y.shape[0])
-    n_points = n_points - n_points % 1024
-    x, y = x[:n_points], y[:n_points]
-    if x.ndim == 2 and y.ndim == 2:
-        x, y = x[None], y[None]
-    x, y = torch.from_numpy(x).cuda(), torch.from_numpy(y).cuda()
-    dist, _ = module(x, y, 0.005, 50)
-    dist = torch.sqrt(dist).mean()
-    return dist.item()

lidm/eval/models/minkowskinet/model.py

@@ -1,141 +0,0 @@
-import torch
-import torch.nn as nn
-
-try:
-    import torchsparse
-    import torchsparse.nn as spnn
-    from ..ts import basic_blocks
-except ImportError:
-    raise Exception('Required ts lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
-
-
-class Model(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-
-        cr = config.model_params.cr
-        cs = config.model_params.layer_num
-        cs = [int(cr * x) for x in cs]
-
-        self.pres = self.vres = config.model_params.voxel_size
-        self.num_classes = config.model_params.num_class
-
-        self.stem = nn.Sequential(
-            spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
-            spnn.BatchNorm(cs[0]), spnn.ReLU(True),
-            spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
-            spnn.BatchNorm(cs[0]), spnn.ReLU(True))
-
-        self.stage1 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage2 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage3 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage4 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
-        )
-
-        self.up1 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up2 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up3 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up4 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
-
-        self.weight_initialization()
-        self.dropout = nn.Dropout(0.3, True)
-
-    def weight_initialization(self):
-        for m in self.modules():
-            if isinstance(m, nn.BatchNorm1d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-
-    def forward(self, data_dict, return_logits=False, return_final_logits=False):
-        x = data_dict['lidar']
-        x.C = x.C.int()
-
-        x0 = self.stem(x)
-        x1 = self.stage1(x0)
-        x2 = self.stage2(x1)
-        x3 = self.stage3(x2)
-        x4 = self.stage4(x3)
-
-        if return_logits:
-            output_dict = dict()
-            output_dict['logits'] = x4.F
-            output_dict['batch_indices'] = x4.C[:, -1]
-            return output_dict
-
-        y1 = self.up1[0](x4)
-        y1 = torchsparse.cat([y1, x3])
-        y1 = self.up1[1](y1)
-
-        y2 = self.up2[0](y1)
-        y2 = torchsparse.cat([y2, x2])
-        y2 = self.up2[1](y2)
-
-        y3 = self.up3[0](y2)
-        y3 = torchsparse.cat([y3, x1])
-        y3 = self.up3[1](y3)
-
-        y4 = self.up4[0](y3)
-        y4 = torchsparse.cat([y4, x0])
-        y4 = self.up4[1](y4)
-        if return_final_logits:
-            output_dict = dict()
-            output_dict['logits'] = y4.F
-            output_dict['coords'] = y4.C[:, :3]
-            output_dict['batch_indices'] = y4.C[:, -1]
-            return output_dict
-
-        output = self.classifier(y4.F)
-        data_dict['output'] = output.F
-
-        return data_dict

lidm/eval/models/rangenet/model.py

@@ -1,372 +0,0 @@
-#!/usr/bin/env python3
-# This file is covered by the LICENSE file in the root of this project.
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class BasicBlock(nn.Module):
-    def __init__(self, inplanes, planes, bn_d=0.1):
-        super(BasicBlock, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
-                               stride=1, padding=0, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes[0], momentum=bn_d)
-        self.relu1 = nn.LeakyReLU(0.1)
-        self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
-                               stride=1, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes[1], momentum=bn_d)
-        self.relu2 = nn.LeakyReLU(0.1)
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu1(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu2(out)
-
-        out += residual
-        return out
-
-
-# ******************************************************************************
-
-# number of layers per model
-model_blocks = {
-    21: [1, 1, 2, 2, 1],
-    53: [1, 2, 8, 8, 4],
-}
-
-
-class Backbone(nn.Module):
-    """
-       Class for DarknetSeg. Subclasses PyTorch's own "nn" module
-    """
-
-    def __init__(self, params):
-        super(Backbone, self).__init__()
-        self.use_range = params["input_depth"]["range"]
-        self.use_xyz = params["input_depth"]["xyz"]
-        self.use_remission = params["input_depth"]["remission"]
-        self.drop_prob = params["dropout"]
-        self.bn_d = params["bn_d"]
-        self.OS = params["OS"]
-        self.layers = params["extra"]["layers"]
-
-        # input depth calc
-        self.input_depth = 0
-        self.input_idxs = []
-        if self.use_range:
-            self.input_depth += 1
-            self.input_idxs.append(0)
-        if self.use_xyz:
-            self.input_depth += 3
-            self.input_idxs.extend([1, 2, 3])
-        if self.use_remission:
-            self.input_depth += 1
-            self.input_idxs.append(4)
-
-        # stride play
-        self.strides = [2, 2, 2, 2, 2]
-        # check current stride
-        current_os = 1
-        for s in self.strides:
-            current_os *= s
-
-        # make the new stride
-        if self.OS > current_os:
-            print("Can't do OS, ", self.OS,
-                  " because it is bigger than original ", current_os)
-        else:
-            # redo strides according to needed stride
-            for i, stride in enumerate(reversed(self.strides), 0):
-                if int(current_os) != self.OS:
-                    if stride == 2:
-                        current_os /= 2
-                        self.strides[-1 - i] = 1
-                    if int(current_os) == self.OS:
-                        break
-
-        # check that darknet exists
-        assert self.layers in model_blocks.keys()
-
-        # generate layers depending on darknet type
-        self.blocks = model_blocks[self.layers]
-
-        # input layer
-        self.conv1 = nn.Conv2d(self.input_depth, 32, kernel_size=3,
-                               stride=1, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(32, momentum=self.bn_d)
-        self.relu1 = nn.LeakyReLU(0.1)
-
-        # encoder
-        self.enc1 = self._make_enc_layer(BasicBlock, [32, 64], self.blocks[0],
-                                         stride=self.strides[0], bn_d=self.bn_d)
-        self.enc2 = self._make_enc_layer(BasicBlock, [64, 128], self.blocks[1],
-                                         stride=self.strides[1], bn_d=self.bn_d)
-        self.enc3 = self._make_enc_layer(BasicBlock, [128, 256], self.blocks[2],
-                                         stride=self.strides[2], bn_d=self.bn_d)
-        self.enc4 = self._make_enc_layer(BasicBlock, [256, 512], self.blocks[3],
-                                         stride=self.strides[3], bn_d=self.bn_d)
-        self.enc5 = self._make_enc_layer(BasicBlock, [512, 1024], self.blocks[4],
-                                         stride=self.strides[4], bn_d=self.bn_d)
-
-        # for a bit of fun
-        self.dropout = nn.Dropout2d(self.drop_prob)
-
-        # last channels
-        self.last_channels = 1024
-
-    # make layer useful function
-    def _make_enc_layer(self, block, planes, blocks, stride, bn_d=0.1):
-        layers = []
-
-        #  downsample
-        layers.append(("conv", nn.Conv2d(planes[0], planes[1],
-                                         kernel_size=3,
-                                         stride=[1, stride], dilation=1,
-                                         padding=1, bias=False)))
-        layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
-        layers.append(("relu", nn.LeakyReLU(0.1)))
-
-        #  blocks
-        inplanes = planes[1]
-        for i in range(0, blocks):
-            layers.append(("residual_{}".format(i),
-                           block(inplanes, planes, bn_d)))
-
-        return nn.Sequential(OrderedDict(layers))
-
-    def run_layer(self, x, layer, skips, os):
-        y = layer(x)
-        if y.shape[2] < x.shape[2] or y.shape[3] < x.shape[3]:
-            skips[os] = x.detach()
-            os *= 2
-        x = y
-        return x, skips, os
-
-    def forward(self, x, return_logits=False, return_list=None):
-        # filter input
-        x = x[:, self.input_idxs]
-
-        # run cnn
-        # store for skip connections
-        skips = {}
-        out_dict = {}
-        os = 1
-
-        # first layer
-        x, skips, os = self.run_layer(x, self.conv1, skips, os)
-        x, skips, os = self.run_layer(x, self.bn1, skips, os)
-        x, skips, os = self.run_layer(x, self.relu1, skips, os)
-        if return_list and 'enc_0' in return_list:
-            out_dict['enc_0'] = x.detach().cpu()  # 32, 64, 1024
-
-        # all encoder blocks with intermediate dropouts
-        x, skips, os = self.run_layer(x, self.enc1, skips, os)
-        if return_list and 'enc_1' in return_list:
-            out_dict['enc_1'] = x.detach().cpu()  # 64, 64, 512
-        x, skips, os = self.run_layer(x, self.dropout, skips, os)
-
-        x, skips, os = self.run_layer(x, self.enc2, skips, os)
-        if return_list and 'enc_2' in return_list:
-            out_dict['enc_2'] = x.detach().cpu()  # 128, 64, 256
-        x, skips, os = self.run_layer(x, self.dropout, skips, os)
-
-        x, skips, os = self.run_layer(x, self.enc3, skips, os)
-        if return_list and 'enc_3' in return_list:
-            out_dict['enc_3'] = x.detach().cpu()  # 256, 64, 128
-        x, skips, os = self.run_layer(x, self.dropout, skips, os)
-
-        x, skips, os = self.run_layer(x, self.enc4, skips, os)
-        if return_list and 'enc_4' in return_list:
-            out_dict['enc_4'] = x.detach().cpu()  # 512, 64, 64
-        x, skips, os = self.run_layer(x, self.dropout, skips, os)
-
-        x, skips, os = self.run_layer(x, self.enc5, skips, os)
-        if return_list and 'enc_5' in return_list:
-            out_dict['enc_5'] = x.detach().cpu()  # 1024, 64, 32
-        if return_logits:
-            return x
-
-        x, skips, os = self.run_layer(x, self.dropout, skips, os)
-
-        if return_list is not None:
-            return x, skips, out_dict
-        return x, skips
-
-    def get_last_depth(self):
-        return self.last_channels
-
-    def get_input_depth(self):
-        return self.input_depth
-
-
-class Decoder(nn.Module):
-    """
-       Class for DarknetSeg. Subclasses PyTorch's own "nn" module
-    """
-
-    def __init__(self, params, OS=32, feature_depth=1024):
-        super(Decoder, self).__init__()
-        self.backbone_OS = OS
-        self.backbone_feature_depth = feature_depth
-        self.drop_prob = params["dropout"]
-        self.bn_d = params["bn_d"]
-        self.index = 0
-
-        # stride play
-        self.strides = [2, 2, 2, 2, 2]
-        # check current stride
-        current_os = 1
-        for s in self.strides:
-            current_os *= s
-        # redo strides according to needed stride
-        for i, stride in enumerate(self.strides):
-            if int(current_os) != self.backbone_OS:
-                if stride == 2:
-                    current_os /= 2
-                    self.strides[i] = 1
-                if int(current_os) == self.backbone_OS:
-                    break
-
-        # decoder
-        self.dec5 = self._make_dec_layer(BasicBlock,
-                                         [self.backbone_feature_depth, 512],
-                                         bn_d=self.bn_d,
-                                         stride=self.strides[0])
-        self.dec4 = self._make_dec_layer(BasicBlock, [512, 256], bn_d=self.bn_d,
-                                         stride=self.strides[1])
-        self.dec3 = self._make_dec_layer(BasicBlock, [256, 128], bn_d=self.bn_d,
-                                         stride=self.strides[2])
-        self.dec2 = self._make_dec_layer(BasicBlock, [128, 64], bn_d=self.bn_d,
-                                         stride=self.strides[3])
-        self.dec1 = self._make_dec_layer(BasicBlock, [64, 32], bn_d=self.bn_d,
-                                         stride=self.strides[4])
-
-        # layer list to execute with skips
-        self.layers = [self.dec5, self.dec4, self.dec3, self.dec2, self.dec1]
-
-        # for a bit of fun
-        self.dropout = nn.Dropout2d(self.drop_prob)
-
-        # last channels
-        self.last_channels = 32
-
-    def _make_dec_layer(self, block, planes, bn_d=0.1, stride=2):
-        layers = []
-
-        #  downsample
-        if stride == 2:
-            layers.append(("upconv", nn.ConvTranspose2d(planes[0], planes[1],
-                                                        kernel_size=[1, 4], stride=[1, 2],
-                                                        padding=[0, 1])))
-        else:
-            layers.append(("conv", nn.Conv2d(planes[0], planes[1],
-                                             kernel_size=3, padding=1)))
-        layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
-        layers.append(("relu", nn.LeakyReLU(0.1)))
-
-        #  blocks
-        layers.append(("residual", block(planes[1], planes, bn_d)))
-
-        return nn.Sequential(OrderedDict(layers))
-
-    def run_layer(self, x, layer, skips, os):
-        feats = layer(x)  # up
-        if feats.shape[-1] > x.shape[-1]:
-            os //= 2  # match skip
-            feats = feats + skips[os].detach()  # add skip
-        x = feats
-        return x, skips, os
-
-    def forward(self, x, skips, return_logits=False, return_list=None):
-        os = self.backbone_OS
-        out_dict = {}
-
-        # run layers
-        x, skips, os = self.run_layer(x, self.dec5, skips, os)
-        if return_list and 'dec_4' in return_list:
-            out_dict['dec_4'] = x.detach().cpu()  # 512, 64, 64
-        x, skips, os = self.run_layer(x, self.dec4, skips, os)
-        if return_list and 'dec_3' in return_list:
-            out_dict['dec_3'] = x.detach().cpu()  # 256, 64, 128
-        x, skips, os = self.run_layer(x, self.dec3, skips, os)
-        if return_list and 'dec_2' in return_list:
-            out_dict['dec_2'] = x.detach().cpu()  # 128, 64, 256
-        x, skips, os = self.run_layer(x, self.dec2, skips, os)
-        if return_list and 'dec_1' in return_list:
-            out_dict['dec_1'] = x.detach().cpu()  # 64, 64, 512
-        x, skips, os = self.run_layer(x, self.dec1, skips, os)
-        if return_list and 'dec_0' in return_list:
-            out_dict['dec_0'] = x.detach().cpu()  # 32, 64, 1024
-
-        logits = torch.clone(x).detach()
-        x = self.dropout(x)
-
-        if return_logits:
-            return x, logits
-        if return_list is not None:
-            return out_dict
-        return x
-
-    def get_last_depth(self):
-        return self.last_channels
-
-
-class Model(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.backbone = Backbone(params=self.config["backbone"])
-        self.decoder = Decoder(params=self.config["decoder"], OS=self.config["backbone"]["OS"],
-                               feature_depth=self.backbone.get_last_depth())
-
-    def load_pretrained_weights(self, path):
-        w_dict = torch.load(path + "/backbone",
-                            map_location=lambda storage, loc: storage)
-        self.backbone.load_state_dict(w_dict, strict=True)
-        w_dict = torch.load(path + "/segmentation_decoder",
-                            map_location=lambda storage, loc: storage)
-        self.decoder.load_state_dict(w_dict, strict=True)
-
-    def forward(self, x, return_logits=False, return_final_logits=False, return_list=None, agg_type='depth'):
-        if return_logits:
-            logits = self.backbone(x, return_logits)
-            logits = F.adaptive_avg_pool2d(logits, (1, 1)).squeeze()
-            logits = torch.clone(logits).detach().cpu().numpy()
-            return logits
-        elif return_list is not None:
-            x, skips, enc_dict = self.backbone(x, return_list=return_list)
-            dec_dict = self.decoder(x, skips, return_list=return_list)
-            out_dict = {**enc_dict, **dec_dict}
-            return out_dict
-        elif return_final_logits:
-            assert agg_type in ['all', 'sector', 'depth']
-            y, skips = self.backbone(x)
-            y, logits = self.decoder(y, skips, True)
-
-            B, C, H, W = logits.shape
-            N = 16
-
-            # avg all
-            if agg_type == 'all':
-                logits = logits.mean([2, 3])
-            # avg in patch
-            elif agg_type == 'sector':
-                logits = logits.view(B, C, H, N, W // N).mean([2, 4]).reshape(B, -1)
-            # avg in row
-            elif agg_type == 'depth':
-                logits = logits.view(B, C, N, H // N, W).mean([3, 4]).reshape(B, -1)
-
-            logits = torch.clone(logits).detach().cpu().numpy()
-            return logits
-        else:
-            y, skips = self.backbone(x)
-            y = self.decoder(y, skips, False)
-            return y

lidm/eval/models/spvcnn/model.py

@@ -1,179 +0,0 @@
-import torch.nn as nn
-
-try:
-    import torchsparse
-    import torchsparse.nn as spnn
-    from torchsparse import PointTensor
-    from ..ts.utils import initial_voxelize, point_to_voxel, voxel_to_point
-    from ..ts import basic_blocks
-except ImportError:
-    raise Exception('Required torchsparse lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
-
-
-class Model(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        cr = config.model_params.cr
-        cs = config.model_params.layer_num
-        cs = [int(cr * x) for x in cs]
-
-        self.pres = self.vres = config.model_params.voxel_size
-        self.num_classes = config.model_params.num_class
-
-        self.stem = nn.Sequential(
-            spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
-            spnn.BatchNorm(cs[0]), spnn.ReLU(True),
-            spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
-            spnn.BatchNorm(cs[0]), spnn.ReLU(True))
-
-        self.stage1 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage2 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage3 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
-        )
-
-        self.stage4 = nn.Sequential(
-            basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
-            basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
-            basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
-        )
-
-        self.up1 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up2 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up3 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.up4 = nn.ModuleList([
-            basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
-            nn.Sequential(
-                basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
-                                           dilation=1),
-                basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
-            )
-        ])
-
-        self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
-
-        self.point_transforms = nn.ModuleList([
-            nn.Sequential(
-                nn.Linear(cs[0], cs[4]),
-                nn.BatchNorm1d(cs[4]),
-                nn.ReLU(True),
-            ),
-            nn.Sequential(
-                nn.Linear(cs[4], cs[6]),
-                nn.BatchNorm1d(cs[6]),
-                nn.ReLU(True),
-            ),
-            nn.Sequential(
-                nn.Linear(cs[6], cs[8]),
-                nn.BatchNorm1d(cs[8]),
-                nn.ReLU(True),
-            )
-        ])
-
-        self.weight_initialization()
-        self.dropout = nn.Dropout(0.3, True)
-
-    def weight_initialization(self):
-        for m in self.modules():
-            if isinstance(m, nn.BatchNorm1d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-
-    def forward(self, data_dict, return_logits=False, return_final_logits=False):
-        x = data_dict['lidar']
-
-        # x: SparseTensor z: PointTensor
-        z = PointTensor(x.F, x.C.float())
-
-        x0 = initial_voxelize(z, self.pres, self.vres)
-
-        x0 = self.stem(x0)
-        z0 = voxel_to_point(x0, z, nearest=False)
-        z0.F = z0.F
-
-        x1 = point_to_voxel(x0, z0)
-        x1 = self.stage1(x1)
-        x2 = self.stage2(x1)
-        x3 = self.stage3(x2)
-        x4 = self.stage4(x3)
-        z1 = voxel_to_point(x4, z0)
-        z1.F = z1.F + self.point_transforms[0](z0.F)
-
-        y1 = point_to_voxel(x4, z1)
-
-        if return_logits:
-            output_dict = dict()
-            output_dict['logits'] = y1.F
-            output_dict['batch_indices'] = y1.C[:, -1]
-            return output_dict
-
-        y1.F = self.dropout(y1.F)
-        y1 = self.up1[0](y1)
-        y1 = torchsparse.cat([y1, x3])
-        y1 = self.up1[1](y1)
-
-        y2 = self.up2[0](y1)
-        y2 = torchsparse.cat([y2, x2])
-        y2 = self.up2[1](y2)
-        z2 = voxel_to_point(y2, z1)
-        z2.F = z2.F + self.point_transforms[1](z1.F)
-
-        y3 = point_to_voxel(y2, z2)
-        y3.F = self.dropout(y3.F)
-        y3 = self.up3[0](y3)
-        y3 = torchsparse.cat([y3, x1])
-        y3 = self.up3[1](y3)
-
-        y4 = self.up4[0](y3)
-        y4 = torchsparse.cat([y4, x0])
-        y4 = self.up4[1](y4)
-        z3 = voxel_to_point(y4, z2)
-        z3.F = z3.F + self.point_transforms[2](z2.F)
-
-        if return_final_logits:
-            output_dict = dict()
-            output_dict['logits'] = z3.F
-            output_dict['coords'] = z3.C[:, :3]
-            output_dict['batch_indices'] = z3.C[:, -1].long()
-            return output_dict
-
-        # output = self.classifier(z3.F)
-        data_dict['logits'] = z3.F
-
-        return data_dict

lidm/eval/models/ts/basic_blocks.py

@@ -1,79 +0,0 @@
-#!/usr/bin/env python
-# encoding: utf-8
-'''
-@author: Xu Yan
-@file: basic_blocks.py
-@time: 2021/4/14 22:53
-'''
-import torch.nn as nn
-
-try:
-    import torchsparse.nn as spnn
-except:
-    print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
-
-
-class BasicConvolutionBlock(nn.Module):
-    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
-        super().__init__()
-        self.net = nn.Sequential(
-            spnn.Conv3d(
-                inc,
-                outc,
-                kernel_size=ks,
-                dilation=dilation,
-                stride=stride), spnn.BatchNorm(outc),
-            spnn.ReLU(True))
-
-    def forward(self, x):
-        out = self.net(x)
-        return out
-
-
-class BasicDeconvolutionBlock(nn.Module):
-    def __init__(self, inc, outc, ks=3, stride=1):
-        super().__init__()
-        self.net = nn.Sequential(
-            spnn.Conv3d(
-                inc,
-                outc,
-                kernel_size=ks,
-                stride=stride,
-                transposed=True),
-            spnn.BatchNorm(outc),
-            spnn.ReLU(True))
-
-    def forward(self, x):
-        return self.net(x)
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
-        super().__init__()
-        self.net = nn.Sequential(
-            spnn.Conv3d(
-                inc,
-                outc,
-                kernel_size=ks,
-                dilation=dilation,
-                stride=stride), spnn.BatchNorm(outc),
-            spnn.ReLU(True),
-            spnn.Conv3d(
-                outc,
-                outc,
-                kernel_size=ks,
-                dilation=dilation,
-                stride=1),
-            spnn.BatchNorm(outc))
-
-        self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \
-            nn.Sequential(
-                spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride),
-                spnn.BatchNorm(outc)
-            )
-
-        self.ReLU = spnn.ReLU(True)
-
-    def forward(self, x):
-        out = self.ReLU(self.net(x) + self.downsample(x))
-        return out

lidm/eval/models/ts/utils.py

@@ -1,90 +0,0 @@
-import torch
-
-try:
-    import torchsparse.nn.functional as F
-    from torchsparse import PointTensor, SparseTensor
-    from torchsparse.nn.utils import get_kernel_offsets
-except:
-    print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
-
-__all__ = ['initial_voxelize', 'point_to_voxel', 'voxel_to_point']
-
-
-# z: PointTensor
-# return: SparseTensor
-def initial_voxelize(z, init_res, after_res):
-    new_float_coord = torch.cat([(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)
-
-    pc_hash = F.sphash(torch.floor(new_float_coord).int())
-    sparse_hash = torch.unique(pc_hash)
-    idx_query = F.sphashquery(pc_hash, sparse_hash)
-    counts = F.spcount(idx_query.int(), len(sparse_hash))
-
-    inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query, counts)
-    inserted_coords = torch.round(inserted_coords).int()
-    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
-
-    new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
-    new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
-    z.additional_features['idx_query'][1] = idx_query
-    z.additional_features['counts'][1] = counts
-    z.C = new_float_coord
-
-    return new_tensor
-
-
-# x: SparseTensor, z: PointTensor
-# return: SparseTensor
-def point_to_voxel(x, z):
-    if z.additional_features is None or \
-            z.additional_features.get('idx_query') is None or \
-            z.additional_features['idx_query'].get(x.s) is None:
-        pc_hash = F.sphash(
-            torch.cat([torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0], z.C[:, -1].int().view(-1, 1)], 1))
-        sparse_hash = F.sphash(x.C)
-        idx_query = F.sphashquery(pc_hash, sparse_hash)
-        counts = F.spcount(idx_query.int(), x.C.shape[0])
-        z.additional_features['idx_query'][x.s] = idx_query
-        z.additional_features['counts'][x.s] = counts
-    else:
-        idx_query = z.additional_features['idx_query'][x.s]
-        counts = z.additional_features['counts'][x.s]
-
-    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
-    new_tensor = SparseTensor(inserted_feat, x.C, x.s)
-    new_tensor.cmaps = x.cmaps
-    new_tensor.kmaps = x.kmaps
-
-    return new_tensor
-
-
-# x: SparseTensor, z: PointTensor
-# return: PointTensor
-def voxel_to_point(x, z, nearest=False):
-    if z.idx_query is None or z.weights is None or z.idx_query.get(x.s) is None or z.weights.get(x.s) is None:
-        off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
-        old_hash = F.sphash(
-            torch.cat([
-                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
-                z.C[:, -1].int().view(-1, 1)], 1), off)
-        pc_hash = F.sphash(x.C.to(z.F.device))
-        idx_query = F.sphashquery(old_hash, pc_hash)
-        weights = F.calc_ti_weights(z.C, idx_query, scale=x.s[0]).transpose(0, 1).contiguous()
-        idx_query = idx_query.transpose(0, 1).contiguous()
-        if nearest:
-            weights[:, 1:] = 0.
-            idx_query[:, 1:] = -1
-        new_feat = F.spdevoxelize(x.F, idx_query, weights)
-        new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
-        new_tensor.additional_features = z.additional_features
-        new_tensor.idx_query[x.s] = idx_query
-        new_tensor.weights[x.s] = weights
-        z.idx_query[x.s] = idx_query
-        z.weights[x.s] = weights
-
-    else:
-        new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s), z.weights.get(x.s))
-        new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
-        new_tensor.additional_features = z.additional_features
-
-    return new_tensor

lidm/eval/modules/chamfer2D/chamfer2D.cu

@@ -1,182 +0,0 @@
-
-#include <stdio.h>
-#include <ATen/ATen.h>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <vector>
-
-
-
-__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
-	const int batch=512;
-	__shared__ float buf[batch*2];
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int k2=0;k2<m;k2+=batch){
-			int end_k=min(m,k2+batch)-k2;
-			for (int j=threadIdx.x;j<end_k*2;j+=blockDim.x){
-				buf[j]=xyz2[(i*m+k2)*2+j];
-			}
-			__syncthreads();
-			for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-				float x1=xyz[(i*n+j)*2+0];
-				float y1=xyz[(i*n+j)*2+1];
-				int best_i=0;
-				float best=0;
-				int end_ka=end_k-(end_k&2);
-				if (end_ka==batch){
-					for (int k=0;k<batch;k+=4){
-						{
-							float x2=buf[k*2+0]-x1;
-							float y2=buf[k*2+1]-y1;
-							float d=x2*x2+y2*y2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*2+2]-x1;
-							float y2=buf[k*2+3]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*2+4]-x1;
-							float y2=buf[k*2+5]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*2+6]-x1;
-							float y2=buf[k*2+7]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}else{
-					for (int k=0;k<end_ka;k+=4){
-						{
-							float x2=buf[k*2+0]-x1;
-							float y2=buf[k*2+1]-y1;
-							float d=x2*x2+y2*y2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*2+2]-x1;
-							float y2=buf[k*2+3]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*2+4]-x1;
-							float y2=buf[k*2+5]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*2+6]-x1;
-							float y2=buf[k*2+7]-y1;
-							float d=x2*x2+y2*y2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}
-				for (int k=end_ka;k<end_k;k++){
-					float x2=buf[k*2+0]-x1;
-					float y2=buf[k*2+1]-y1;
-					float d=x2*x2+y2*y2;
-					if (k==0 || d<best){
-						best=d;
-						best_i=k+k2;
-					}
-				}
-				if (k2==0 || result[(i*n+j)]>best){
-					result[(i*n+j)]=best;
-					result_i[(i*n+j)]=best_i;
-				}
-			}
-			__syncthreads();
-		}
-	}
-}
-// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){
-
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, n, xyz1.data<float>(), m, xyz2.data<float>(), dist1.data<float>(), idx1.data<int>());
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, m, xyz2.data<float>(), n, xyz1.data<float>(), dist2.data<float>(), idx2.data<int>());
-
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-
-
-}
-__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-			float x1=xyz1[(i*n+j)*2+0];
-			float y1=xyz1[(i*n+j)*2+1];
-			int j2=idx1[i*n+j];
-			float x2=xyz2[(i*m+j2)*2+0];
-			float y2=xyz2[(i*m+j2)*2+1];
-			float g=grad_dist1[i*n+j]*2;
-			atomicAdd(&(grad_xyz1[(i*n+j)*2+0]),g*(x1-x2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*2+1]),g*(y1-y2));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*2+0]),-(g*(x1-x2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*2+1]),-(g*(y1-y2)));
-		}
-	}
-}
-// int chamfer_cuda_backward(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream){
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2){
-	// cudaMemset(grad_xyz1,0,b*n*3*4);
-	// cudaMemset(grad_xyz2,0,b*m*3*4);
-	
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,n,xyz1.data<float>(),m,xyz2.data<float>(),graddist1.data<float>(),idx1.data<int>(),gradxyz1.data<float>(),gradxyz2.data<float>());
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,m,xyz2.data<float>(),n,xyz1.data<float>(),graddist2.data<float>(),idx2.data<int>(),gradxyz2.data<float>(),gradxyz1.data<float>());
-	
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-	
-}
-

lidm/eval/modules/chamfer2D/chamfer_cuda.cpp

@@ -1,33 +0,0 @@
-#include <torch/torch.h>
-#include <vector>
-
-///TMP
-//#include "common.h"
-/// NOT TMP
-	
-
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2);
-
-
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2);
-
-
-
-
-int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) {
-    return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2);
-}
-
-
-int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, 
-					  at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) {
-
-    return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2);
-}
-
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("forward", &chamfer_forward, "chamfer forward (CUDA)");
-  m.def("backward", &chamfer_backward, "chamfer backward (CUDA)");
-}

lidm/eval/modules/chamfer2D/dist_chamfer_2D.py

@@ -1,84 +0,0 @@
-from torch import nn
-from torch.autograd import Function
-import torch
-import importlib
-import os
-
-chamfer_found = importlib.find_loader("chamfer_2D") is not None
-if not chamfer_found:
-    ## Cool trick from https://github.com/chrdiller
-    print("Jitting Chamfer 2D")
-    cur_path = os.path.dirname(os.path.abspath(__file__))
-    build_path = cur_path.replace('chamfer2D', 'tmp')
-    os.makedirs(build_path, exist_ok=True)
-
-    from torch.utils.cpp_extension import load
-
-    chamfer_2D = load(name="chamfer_2D",
-                      sources=[
-                          "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]),
-                          "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer2D.cu"]),
-                      ], build_directory=build_path)
-    print("Loaded JIT 2D CUDA chamfer distance")
-
-else:
-    import chamfer_2D
-
-    print("Loaded compiled 2D CUDA chamfer distance")
-
-
-# Chamfer's distance module @thibaultgroueix
-# GPU tensors only
-class chamfer_2DFunction(Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2):
-        batchsize, n, dim = xyz1.size()
-        assert dim == 2, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
-        _, m, dim = xyz2.size()
-        assert dim == 2, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
-        device = xyz1.device
-
-        device = xyz1.device
-
-        dist1 = torch.zeros(batchsize, n)
-        dist2 = torch.zeros(batchsize, m)
-
-        idx1 = torch.zeros(batchsize, n).type(torch.IntTensor)
-        idx2 = torch.zeros(batchsize, m).type(torch.IntTensor)
-
-        dist1 = dist1.to(device)
-        dist2 = dist2.to(device)
-        idx1 = idx1.to(device)
-        idx2 = idx2.to(device)
-        torch.cuda.set_device(device)
-
-        chamfer_2D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
-        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
-        return dist1, dist2, idx1, idx2
-
-    @staticmethod
-    def backward(ctx, graddist1, graddist2, gradidx1, gradidx2):
-        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-        graddist1 = graddist1.contiguous()
-        graddist2 = graddist2.contiguous()
-        device = graddist1.device
-
-        gradxyz1 = torch.zeros(xyz1.size())
-        gradxyz2 = torch.zeros(xyz2.size())
-
-        gradxyz1 = gradxyz1.to(device)
-        gradxyz2 = gradxyz2.to(device)
-        chamfer_2D.backward(
-            xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
-        )
-        return gradxyz1, gradxyz2
-
-
-class chamfer_2DDist(nn.Module):
-    def __init__(self):
-        super(chamfer_2DDist, self).__init__()
-
-    def forward(self, input1, input2):
-        input1 = input1.contiguous()
-        input2 = input2.contiguous()
-        return chamfer_2DFunction.apply(input1, input2)

lidm/eval/modules/chamfer2D/setup.py

@@ -1,14 +0,0 @@
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='chamfer_2D',
-    ext_modules=[
-        CUDAExtension('chamfer_2D', [
-            "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
-            "/".join(__file__.split('/')[:-1] + ['chamfer2D.cu']),
-        ]),
-    ],
-    cmdclass={
-        'build_ext': BuildExtension
-    })

lidm/eval/modules/chamfer3D/chamfer3D.cu

@@ -1,196 +0,0 @@
-
-#include <stdio.h>
-#include <ATen/ATen.h>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <vector>
-
-
-
-__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
-	const int batch=512;
-	__shared__ float buf[batch*3];
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int k2=0;k2<m;k2+=batch){
-			int end_k=min(m,k2+batch)-k2;
-			for (int j=threadIdx.x;j<end_k*3;j+=blockDim.x){
-				buf[j]=xyz2[(i*m+k2)*3+j];
-			}
-			__syncthreads();
-			for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-				float x1=xyz[(i*n+j)*3+0];
-				float y1=xyz[(i*n+j)*3+1];
-				float z1=xyz[(i*n+j)*3+2];
-				int best_i=0;
-				float best=0;
-				int end_ka=end_k-(end_k&3);
-				if (end_ka==batch){
-					for (int k=0;k<batch;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}else{
-					for (int k=0;k<end_ka;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}
-				for (int k=end_ka;k<end_k;k++){
-					float x2=buf[k*3+0]-x1;
-					float y2=buf[k*3+1]-y1;
-					float z2=buf[k*3+2]-z1;
-					float d=x2*x2+y2*y2+z2*z2;
-					if (k==0 || d<best){
-						best=d;
-						best_i=k+k2;
-					}
-				}
-				if (k2==0 || result[(i*n+j)]>best){
-					result[(i*n+j)]=best;
-					result_i[(i*n+j)]=best_i;
-				}
-			}
-			__syncthreads();
-		}
-	}
-}
-// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){
-
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, n, xyz1.data<float>(), m, xyz2.data<float>(), dist1.data<float>(), idx1.data<int>());
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, m, xyz2.data<float>(), n, xyz1.data<float>(), dist2.data<float>(), idx2.data<int>());
-
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-
-
-}
-__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-			float x1=xyz1[(i*n+j)*3+0];
-			float y1=xyz1[(i*n+j)*3+1];
-			float z1=xyz1[(i*n+j)*3+2];
-			int j2=idx1[i*n+j];
-			float x2=xyz2[(i*m+j2)*3+0];
-			float y2=xyz2[(i*m+j2)*3+1];
-			float z2=xyz2[(i*m+j2)*3+2];
-			float g=grad_dist1[i*n+j]*2;
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+0]),g*(x1-x2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+1]),g*(y1-y2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+2]),g*(z1-z2));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+0]),-(g*(x1-x2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+1]),-(g*(y1-y2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+2]),-(g*(z1-z2)));
-		}
-	}
-}
-// int chamfer_cuda_backward(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream){
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2){
-	// cudaMemset(grad_xyz1,0,b*n*3*4);
-	// cudaMemset(grad_xyz2,0,b*m*3*4);
-	
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,n,xyz1.data<float>(),m,xyz2.data<float>(),graddist1.data<float>(),idx1.data<int>(),gradxyz1.data<float>(),gradxyz2.data<float>());
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,m,xyz2.data<float>(),n,xyz1.data<float>(),graddist2.data<float>(),idx2.data<int>(),gradxyz2.data<float>(),gradxyz1.data<float>());
-	
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-	
-}
-

lidm/eval/modules/chamfer3D/chamfer_cuda.cpp

@@ -1,33 +0,0 @@
-#include <torch/torch.h>
-#include <vector>
-
-///TMP
-//#include "common.h"
-/// NOT TMP
-	
-
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2);
-
-
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2);
-
-
-
-
-int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) {
-    return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2);
-}
-
-
-int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, 
-					  at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) {
-
-    return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2);
-}
-
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("forward", &chamfer_forward, "chamfer forward (CUDA)");
-  m.def("backward", &chamfer_backward, "chamfer backward (CUDA)");
-}

lidm/eval/modules/chamfer3D/dist_chamfer_3D.py

@@ -1,76 +0,0 @@
-from torch import nn
-from torch.autograd import Function
-import torch
-import importlib
-import os
-
-chamfer_found = importlib.find_loader("chamfer_3D") is not None
-if not chamfer_found:
-    ## Cool trick from https://github.com/chrdiller
-    print("Jitting Chamfer 3D")
-
-    from torch.utils.cpp_extension import load
-
-    chamfer_3D = load(name="chamfer_3D",
-                      sources=[
-                          "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]),
-                          "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]),
-                      ])
-    print("Loaded JIT 3D CUDA chamfer distance")
-
-else:
-    import chamfer_3D
-    print("Loaded compiled 3D CUDA chamfer distance")
-
-
-# Chamfer's distance module @thibaultgroueix
-# GPU tensors only
-class chamfer_3DFunction(Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2):
-        batchsize, n, _ = xyz1.size()
-        _, m, _ = xyz2.size()
-        device = xyz1.device
-
-        dist1 = torch.zeros(batchsize, n)
-        dist2 = torch.zeros(batchsize, m)
-
-        idx1 = torch.zeros(batchsize, n).type(torch.IntTensor)
-        idx2 = torch.zeros(batchsize, m).type(torch.IntTensor)
-
-        dist1 = dist1.to(device)
-        dist2 = dist2.to(device)
-        idx1 = idx1.to(device)
-        idx2 = idx2.to(device)
-        torch.cuda.set_device(device)
-
-        chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
-        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
-        return dist1, dist2, idx1, idx2
-
-    @staticmethod
-    def backward(ctx, graddist1, graddist2, gradidx1, gradidx2):
-        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-        graddist1 = graddist1.contiguous()
-        graddist2 = graddist2.contiguous()
-        device = graddist1.device
-
-        gradxyz1 = torch.zeros(xyz1.size())
-        gradxyz2 = torch.zeros(xyz2.size())
-
-        gradxyz1 = gradxyz1.to(device)
-        gradxyz2 = gradxyz2.to(device)
-        chamfer_3D.backward(
-            xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
-        )
-        return gradxyz1, gradxyz2
-
-
-class chamfer_3DDist(nn.Module):
-    def __init__(self):
-        super(chamfer_3DDist, self).__init__()
-
-    def forward(self, input1, input2):
-        input1 = input1.contiguous()
-        input2 = input2.contiguous()
-        return chamfer_3DFunction.apply(input1, input2)

lidm/eval/modules/chamfer3D/setup.py

@@ -1,14 +0,0 @@
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='chamfer_3D',
-    ext_modules=[
-        CUDAExtension('chamfer_3D', [
-            "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
-            "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']),
-        ]),
-    ],
-    cmdclass={
-        'build_ext': BuildExtension
-    })

lidm/eval/modules/emd/emd.cpp

@@ -1,31 +0,0 @@
-// EMD approximation module (based on auction algorithm)
-// author: Minghua Liu
-#include <torch/extension.h>
-#include <vector>
-
-int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 
-	                 at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
-	                 at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters);
-
-int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx);
-
-
-
-int emd_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 
-	                 at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
-	                 at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) {
-	return emd_cuda_forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters);
-}
-
-int emd_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) {
-
-    return emd_cuda_backward(xyz1, xyz2, gradxyz, graddist, idx);
-}
-
-
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("forward", &emd_forward, "emd forward (CUDA)");
-  m.def("backward", &emd_backward, "emd backward (CUDA)");
-}

lidm/eval/modules/emd/emd_cuda.cu

@@ -1,316 +0,0 @@
-// EMD approximation module (based on auction algorithm)
-// author: Minghua Liu
-#include <stdio.h>
-#include <ATen/ATen.h>
-
-#include <cuda.h>
-#include <iostream>
-#include <cuda_runtime.h>
-
-__device__ __forceinline__ float atomicMax(float *address, float val)
-{
-    int ret = __float_as_int(*address);
-    while(val > __int_as_float(ret))
-    {
-        int old = ret;
-        if((ret = atomicCAS((int *)address, old, __float_as_int(val))) == old)
-            break;
-    }
-    return __int_as_float(ret);
-}
-
-
-__global__ void clear(int b, int * cnt_tmp, int * unass_cnt) {
-	for (int i = threadIdx.x; i < b; i += blockDim.x) {
-		cnt_tmp[i] = 0;
-		unass_cnt[i] = 0;
-	}
-}
-
-__global__ void calc_unass_cnt(int b, int n, int * assignment, int * unass_cnt) { 
-	// count the number of unassigned points in each batch
-	const int BLOCK_SIZE = 1024; 
-	__shared__ int scan_array[BLOCK_SIZE];
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		scan_array[threadIdx.x] = assignment[i * n + blockIdx.y * BLOCK_SIZE + threadIdx.x] == -1 ? 1 : 0;
-		__syncthreads();
-		
-		int stride = 1;
-		while(stride <= BLOCK_SIZE / 2) {
-			int index = (threadIdx.x + 1) * stride * 2 - 1; 
-			if(index < BLOCK_SIZE)
-				scan_array[index] += scan_array[index - stride]; 
-			stride = stride * 2;
-			__syncthreads(); 
-		}
-		__syncthreads();
-		
-		if (threadIdx.x == BLOCK_SIZE - 1) {
-			atomicAdd(&unass_cnt[i], scan_array[threadIdx.x]);
-		}
-		__syncthreads();
-	}
-}
-
-__global__ void calc_unass_cnt_sum(int b, int * unass_cnt, int * unass_cnt_sum) {
-	// count the cumulative sum over over unass_cnt
-	const int BLOCK_SIZE = 512; // batch_size <= 512
-	__shared__ int scan_array[BLOCK_SIZE];
-	scan_array[threadIdx.x] = unass_cnt[threadIdx.x];
-	__syncthreads();
-	
-	int stride = 1;
-	while(stride <= BLOCK_SIZE / 2) {
-		int index = (threadIdx.x + 1) * stride * 2 - 1; 
-		if(index < BLOCK_SIZE)
-			scan_array[index] += scan_array[index - stride]; 
-		stride = stride * 2;
-		__syncthreads(); 
-	}
-	__syncthreads();
-	stride = BLOCK_SIZE / 4; 
-	while(stride > 0) {
-		int index = (threadIdx.x + 1) * stride * 2 - 1; 
-		if((index + stride) < BLOCK_SIZE)
-			scan_array[index + stride] += scan_array[index];
-		stride = stride / 2;
-		__syncthreads(); 
-	}
-	__syncthreads(); 
-	
-	//printf("%d\n", unass_cnt_sum[b - 1]);
-	unass_cnt_sum[threadIdx.x] = scan_array[threadIdx.x];
-}
-
-__global__ void calc_unass_idx(int b, int n, int * assignment, int * unass_idx, int * unass_cnt, int * unass_cnt_sum, int * cnt_tmp) {
-	// list all the unassigned points
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		if (assignment[i * n + blockIdx.y * 1024 + threadIdx.x] == -1) {
-			int idx = atomicAdd(&cnt_tmp[i], 1);
-			unass_idx[unass_cnt_sum[i] - unass_cnt[i] + idx] = blockIdx.y * 1024 + threadIdx.x;
-		} 
-	}
-}
-
-__global__ void Bid(int b, int n, const float * xyz1, const float * xyz2, float eps, int * assignment, int * assignment_inv, float * price, 
-					int * bid, float * bid_increments, float * max_increments, int * unass_cnt, int * unass_cnt_sum, int * unass_idx) {
-	const int batch = 2048, block_size = 1024, block_cnt = n / 1024;
-	__shared__ float xyz2_buf[batch * 3];
-	__shared__ float price_buf[batch];
-	__shared__ float best_buf[block_size];
-	__shared__ float better_buf[block_size];
-	__shared__ int best_i_buf[block_size];
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		int _unass_cnt = unass_cnt[i];
-		if (_unass_cnt == 0)
-			continue;
-		int _unass_cnt_sum = unass_cnt_sum[i];
-		int unass_per_block = (_unass_cnt + block_cnt - 1) / block_cnt;
-		int thread_per_unass = block_size / unass_per_block;
-		int unass_this_block = max(min(_unass_cnt - (int) blockIdx.y * unass_per_block, unass_per_block), 0);
-			
-		float x1, y1, z1, best = -1e9, better = -1e9;
-		int best_i = -1, _unass_id = -1, thread_in_unass;
-
-		if (threadIdx.x < thread_per_unass * unass_this_block) {
-			_unass_id = unass_per_block * blockIdx.y + threadIdx.x / thread_per_unass + _unass_cnt_sum - _unass_cnt;
-			_unass_id = unass_idx[_unass_id];
-			thread_in_unass = threadIdx.x % thread_per_unass;
-
-			x1 = xyz1[(i * n + _unass_id) * 3 + 0];
-			y1 = xyz1[(i * n + _unass_id) * 3 + 1];
-			z1 = xyz1[(i * n + _unass_id) * 3 + 2];
-		}
-
-		for (int k2 = 0; k2 < n; k2 += batch) {
-			int end_k = min(n, k2 + batch) - k2;
-			for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) {
-				xyz2_buf[j] = xyz2[(i * n + k2) * 3 + j];
-			}
-			for (int j = threadIdx.x; j < end_k; j += blockDim.x) {
-				price_buf[j] = price[i * n + k2 + j];
-			}
-			__syncthreads();
-
-			if (_unass_id != -1) {
-				int delta = (end_k + thread_per_unass - 1) / thread_per_unass;
-				int l = thread_in_unass * delta;
-				int r = min((thread_in_unass + 1) * delta, end_k);
-				for (int k = l; k < r; k++) 
-				//if (!last || assignment_inv[i * n + k + k2] == -1)
-				{
-					float x2 = xyz2_buf[k * 3 + 0] - x1;
-					float y2 = xyz2_buf[k * 3 + 1] - y1;
-					float z2 = xyz2_buf[k * 3 + 2] - z1;
-					// the coordinates of points should be normalized to [0, 1]
-					float d = 3.0 - sqrtf(x2 * x2 + y2 * y2 + z2 * z2) - price_buf[k];
-					if (d > best) {
-						better = best;
-						best = d;
-						best_i = k + k2;
-					}
-					else if (d > better) {
-						better = d;
-					}
-				}
-			}
-			__syncthreads();
-		}
-
-		best_buf[threadIdx.x] = best;
-		better_buf[threadIdx.x] = better;
-		best_i_buf[threadIdx.x] = best_i;
-		__syncthreads();
-		
-		if (_unass_id != -1 && thread_in_unass == 0) {
-			for (int j = threadIdx.x + 1; j < threadIdx.x + thread_per_unass; j++) {
-				if (best_buf[j] > best) {
-					better = max(best, better_buf[j]);
-					best = best_buf[j];
-					best_i = best_i_buf[j];
-				}
-				else better = max(better, best_buf[j]);
-			}
-			bid[i * n + _unass_id] = best_i;
-			bid_increments[i * n + _unass_id] = best - better + eps; 
-			atomicMax(&max_increments[i * n + best_i], best - better + eps);
-		}
-	}
-}
-
-__global__ void GetMax(int b, int n, int * assignment, int * bid, float * bid_increments, float * max_increments, int * max_idx) {
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		int j = threadIdx.x + blockIdx.y * blockDim.x;
-		if (assignment[i * n + j] == -1) {
-			int bid_id = bid[i * n + j];
-			float bid_inc = bid_increments[i * n + j];
-			float max_inc = max_increments[i * n + bid_id];
-			if (bid_inc - 1e-6 <= max_inc && max_inc <= bid_inc + 1e-6) 
-			{
-				max_idx[i * n + bid_id] = j;
-			}
-		}
-	}
-}
-
-__global__ void Assign(int b, int n, int * assignment, int * assignment_inv, float * price, int * bid, float * bid_increments, float * max_increments, int * max_idx, bool last) {
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		int j = threadIdx.x + blockIdx.y * blockDim.x;
-		if (assignment[i * n + j] == -1) {
-			int bid_id = bid[i * n + j];
-			if (last || max_idx[i * n + bid_id] == j) 
-			{
-				float bid_inc = bid_increments[i * n + j];
-				int ass_inv = assignment_inv[i * n + bid_id];
-				if (!last && ass_inv != -1) {
-					assignment[i * n + ass_inv] = -1;
-				}
-				assignment_inv[i * n + bid_id] = j;
-				assignment[i * n + j] = bid_id;
-				price[i * n + bid_id] += bid_inc;
-				max_increments[i * n + bid_id] = -1e9;
-			}
-		}
-	}
-}
-
-__global__ void CalcDist(int b, int n, float * xyz1, float * xyz2, float * dist, int * assignment) {
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		int j = threadIdx.x + blockIdx.y * blockDim.x;
-		int k = assignment[i * n + j];
-		float deltax = xyz1[(i * n + j) * 3 + 0] - xyz2[(i * n + k) * 3 + 0];
-		float deltay = xyz1[(i * n + j) * 3 + 1] - xyz2[(i * n + k) * 3 + 1];
-		float deltaz = xyz1[(i * n + j) * 3 + 2] - xyz2[(i * n + k) * 3 + 2];
-		dist[i * n + j] = deltax * deltax + deltay * deltay + deltaz * deltaz;
-	}
-}
-
-int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 
-	                 at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
-	                 at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) {
-
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-	
-	if (n != m) {
-		printf("Input Error! The two point clouds should have the same size.\n");
-		return -1;
-	}
-
-	if (batch_size > 512) {
-		printf("Input Error! The batch size should be less than 512.\n");
-		return -1;
-	}
-
-	if (n % 1024 != 0) {
-		printf("Input Error! The size of the point clouds should be a multiple of 1024.\n");
-		return -1;
-	}
-
-	//cudaEvent_t start,stop;
-	//cudaEventCreate(&start);
-	//cudaEventCreate(&stop);
-	//cudaEventRecord(start);
-	//int iters = 50;
-	for (int i = 0; i < iters; i++) {
-		clear<<<1, batch_size>>>(batch_size, cnt_tmp.data<int>(), unass_cnt.data<int>());
-		calc_unass_cnt<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), unass_cnt.data<int>());
-		calc_unass_cnt_sum<<<1, batch_size>>>(batch_size, unass_cnt.data<int>(), unass_cnt_sum.data<int>());
-		calc_unass_idx<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), unass_idx.data<int>(), unass_cnt.data<int>(), 
-											 unass_cnt_sum.data<int>(), cnt_tmp.data<int>());
-		Bid<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), eps, assignment.data<int>(), assignment_inv.data<int>(), 
-			                          price.data<float>(), bid.data<int>(), bid_increments.data<float>(), max_increments.data<float>(),
-			                          unass_cnt.data<int>(), unass_cnt_sum.data<int>(), unass_idx.data<int>());
-		GetMax<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), bid.data<int>(), bid_increments.data<float>(), max_increments.data<float>(), max_idx.data<int>());
-		Assign<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), assignment_inv.data<int>(), price.data<float>(), bid.data<int>(),
-									  bid_increments.data<float>(), max_increments.data<float>(), max_idx.data<int>(), i == iters - 1);
-	}
-	CalcDist<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), dist.data<float>(), assignment.data<int>());
-	//cudaEventRecord(stop);
-	//cudaEventSynchronize(stop);
-	//float elapsedTime;
-	//cudaEventElapsedTime(&elapsedTime,start,stop);
-	//printf("%lf\n", elapsedTime);
-
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd Output: %s\n", cudaGetErrorString(err));
-	    return 0;
-	  }
-	  return 1;
-}
-
-__global__ void NmDistanceGradKernel(int b, int n, const float * xyz1, const float * xyz2, const float * grad_dist, const int * idx, float * grad_xyz){
-	for (int i = blockIdx.x; i < b; i += gridDim.x) {
-		for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x * gridDim.y) {
-			float x1 = xyz1[(i * n + j) * 3 + 0];
-			float y1 = xyz1[(i * n + j) * 3 + 1];
-			float z1 = xyz1[(i * n + j) * 3 + 2];
-			int j2 = idx[i * n + j];
-			float x2 = xyz2[(i * n + j2) * 3 + 0];
-			float y2 = xyz2[(i * n + j2) * 3 + 1];
-			float z2 = xyz2[(i * n + j2) * 3 + 2];
-			float g = grad_dist[i * n + j] * 2;
-			atomicAdd(&(grad_xyz[(i * n + j) * 3 + 0]), g * (x1 - x2));
-			atomicAdd(&(grad_xyz[(i * n + j) * 3 + 1]), g * (y1 - y2));
-			atomicAdd(&(grad_xyz[(i * n + j) * 3 + 2]), g * (z1 - z2));
-		}
-	}
-}
-
-int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx){
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); 
-	const auto m = xyz2.size(1); 
-
-	NmDistanceGradKernel<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), graddist.data<float>(), idx.data<int>(), gradxyz.data<float>());
-	
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
-	    return 0;
-	  }
-	  return 1;
-	
-}

lidm/eval/modules/emd/emd_module.py

@@ -1,112 +0,0 @@
-# EMD approximation module (based on auction algorithm)
-# memory complexity: O(n)
-# time complexity: O(n^2 * iter) 
-# author: Minghua Liu
-
-# Input:
-# xyz1, xyz2: [#batch, #points, 3]
-# where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud 
-# two point clouds should have same size and be normalized to [0, 1]
-# #points should be a multiple of 1024
-# #batch should be no greater than 512
-# eps is a parameter which balances the error rate and the speed of convergence
-# iters is the number of iteration
-# we only calculate gradient for xyz1
-
-# Output:
-# dist: [#batch, #points],  sqrt(dist) -> L2 distance 
-# assignment: [#batch, #points], index of the matched point in the ground truth point cloud
-# the result is an approximation and the assignment is not guranteed to be a bijection
-import importlib
-import os
-import time
-import numpy as np
-import torch
-from torch import nn
-from torch.autograd import Function
-
-emd_found = importlib.find_loader("emd") is not None
-if not emd_found:
-    ## Cool trick from https://github.com/chrdiller
-    print("Jitting EMD 3D")
-
-    from torch.utils.cpp_extension import load
-
-    emd = load(name="emd",
-               sources=[
-                   "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd.cpp"]),
-                   "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd_cuda.cu"]),
-               ])
-    print("Loaded JIT 3D CUDA emd")
-else:
-    import emd
-    print("Loaded compiled 3D CUDA emd")
-
-
-class emdFunction(Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2, eps, iters):
-        batchsize, n, _ = xyz1.size()
-        _, m, _ = xyz2.size()
-
-        assert (n == m)
-        assert (xyz1.size()[0] == xyz2.size()[0])
-        # assert(n % 1024 == 0)
-        assert (batchsize <= 512)
-
-        xyz1 = xyz1.contiguous().float().cuda()
-        xyz2 = xyz2.contiguous().float().cuda()
-        dist = torch.zeros(batchsize, n, device='cuda').contiguous()
-        assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1
-        assignment_inv = torch.zeros(batchsize, m, device='cuda', dtype=torch.int32).contiguous() - 1
-        price = torch.zeros(batchsize, m, device='cuda').contiguous()
-        bid = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous()
-        bid_increments = torch.zeros(batchsize, n, device='cuda').contiguous()
-        max_increments = torch.zeros(batchsize, m, device='cuda').contiguous()
-        unass_idx = torch.zeros(batchsize * n, device='cuda', dtype=torch.int32).contiguous()
-        max_idx = torch.zeros(batchsize * m, device='cuda', dtype=torch.int32).contiguous()
-        unass_cnt = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
-        unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
-        cnt_tmp = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
-
-        emd.forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx,
-                    unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters)
-
-        ctx.save_for_backward(xyz1, xyz2, assignment)
-        return dist, assignment
-
-    @staticmethod
-    def backward(ctx, graddist, gradidx):
-        xyz1, xyz2, assignment = ctx.saved_tensors
-        graddist = graddist.contiguous()
-
-        gradxyz1 = torch.zeros(xyz1.size(), device='cuda').contiguous()
-        gradxyz2 = torch.zeros(xyz2.size(), device='cuda').contiguous()
-
-        emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment)
-        return gradxyz1, gradxyz2, None, None
-
-
-class emdModule(nn.Module):
-    def __init__(self):
-        super(emdModule, self).__init__()
-
-    def forward(self, input1, input2, eps, iters):
-        return emdFunction.apply(input1, input2, eps, iters)
-
-
-def test_emd():
-    x1 = torch.rand(20, 8192, 3).cuda()
-    x2 = torch.rand(20, 8192, 3).cuda()
-    emd = emdModule()
-    start_time = time.perf_counter()
-    dis, assigment = emd(x1, x2, 0.05, 3000)
-    print("Input_size: ", x1.shape)
-    print("Runtime: %lfs" % (time.perf_counter() - start_time))
-    print("EMD: %lf" % np.sqrt(dis.cpu()).mean())
-    print("|set(assignment)|: %d" % assigment.unique().numel())
-    assigment = assigment.cpu().numpy()
-    assigment = np.expand_dims(assigment, -1)
-    x2 = np.take_along_axis(x2, assigment, axis=1)
-    d = (x1 - x2) * (x1 - x2)
-    print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean())

lidm/eval/modules/emd/setup.py

@@ -1,14 +0,0 @@
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='emd',
-    ext_modules=[
-        CUDAExtension('emd', [
-            'emd.cpp',
-            'emd_cuda.cu',
-        ]),
-    ],
-    cmdclass={
-        'build_ext': BuildExtension
-    })

lidm/models/diffusion/ddim.py

@@ -17,7 +17,7 @@ class DDIMSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
+            if attr.device != torch.device("cuda") and torch.cuda.is_available():
                 attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 

lidm/models/diffusion/plms.py

@@ -17,7 +17,7 @@ class PLMSSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
+            if attr.device != torch.device("cuda") and torch.cuda.is_available():
                 attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 

sample_cond.py

@@ -103,7 +103,7 @@ def make_convolutional_sample(model, batch, batch_size, custom_steps=None, eta=1
 def sample(model, cond):
     batch = {'camera': cond}
     img = make_convolutional_sample(model, batch, batch_size=1, custom_steps=CUSTOM_STEPS, eta=ETA)  # TODO add arguments for batch_size, custom_steps and eta
-    img = img[0, 0]
     pcd = custom_to_pcd(img, model_config)[0].astype(np.float32)
+    img = img.squeeze().detach().cpu().numpy()
     return img, pcd