init

README.md

@@ -1,8 +1,8 @@
 ---
 title: LiDAR Diffusion
-emoji: 📚
-colorFrom: blue
-colorTo: green
+emoji: 🚙🛞🚨
+colorFrom: green
+colorTo: indigo
 sdk: gradio
 sdk_version: 4.26.0
 app_file: app.py

app.py

@@ -1,9 +1,81 @@
 import gradio as gr
+import spaces
+import tempfile
+import os
+import torch
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap
 
+from app_config import CSS, TITLE, DESCRIPTION, DEVICE
+import sample_cond
 
-def greet(name):
-    return "Hello " + name + "!!"
+model = sample_cond.load_model()
 
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+def create_custom_colormap():
+    colors = [(0, 1, 0), (0, 1, 1), (0, 0, 1), (1, 0, 1), (1, 1, 0)]
+    positions = [0, 0.38, 0.6, 0.7, 1]
+
+    custom_cmap = LinearSegmentedColormap.from_list('custom_colormap', list(zip(positions, colors)), N=256)
+    return custom_cmap
+
+
+def colorize_depth(depth, log_scale):
+    if log_scale:
+        depth = ((np.log2((depth / 255.) * 56. + 1) / 5.84) * 255.).astype(np.uint8)
+    mask = depth == 0
+    colormap = create_custom_colormap()
+    rgb = colormap(depth)[:, :, :3]
+    rgb[mask] = 0.
+    return rgb
+
+
+@spaces.GPU
+@torch.no_grad()
+def generate_lidar(model, cond):
+    img, pcd = sample_cond.sample(model, cond)
+    return img, pcd
+
+
+def load_camera(image):
+    split_per_view = 4
+    camera = np.array(image).astype(np.float32) / 255.
+    camera = camera.transpose(2, 0, 1)
+    camera_list = np.split(camera, split_per_view, axis=2)  # split into n chunks as different views
+    camera_cond = torch.from_numpy(np.stack(camera_list, axis=0)).unsqueeze(0).to(DEVICE)
+    return camera_cond
+
+
+with gr.Blocks(css=CSS) as demo:
+    gr.Markdown(TITLE)
+    gr.Markdown(DESCRIPTION)
+    gr.Markdown("### Camera-to-LiDAR Demo")
+    # gr.Markdown("You can slide the output to compare the depth prediction with input image")
+
+    with gr.Row():
+        input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
+        output_image = gr.Image(label="Range Map", elem_id='img-display-output')
+    raw_file = gr.File(label="Point Cloud (.txt file). Can be viewed through Meshlab")
+    submit = gr.Button("Submit")
+
+    def on_submit(image):
+        cond = load_camera(image)
+        img, pcd = generate_lidar(model, cond)
+
+        tmp = tempfile.NamedTemporaryFile(suffix='.txt', delete=False)
+        pcd.save(tmp.name)
+
+        rgb_img = colorize_depth(img, log_scale=True)
+
+        return [rgb_img, tmp.name]
+
+    submit.click(on_submit, inputs=[input_image], outputs=[output_image, raw_file])
+
+    example_files = sorted(os.listdir('cam_examples'))
+    example_files = [os.path.join('cam_examples', filename) for filename in example_files]
+    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[output_image, raw_file],
+                           fn=on_submit, cache_examples=True)
+
+
+if __name__ == '__main__':
+    demo.queue().launch()

app_config.py

@@ -0,0 +1,17 @@
+import torch
+
+CSS = """
+#img-display-container {
+    max-height: 100vh;
+    }
+#img-display-input {
+    max-height: 80vh;
+    }
+#img-display-output {
+    max-height: 80vh;
+    }
+"""
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+TITLE = "# LiDAR Diffusion"
+DESCRIPTION = """Official demo for **LiDAR Diffusion: Towards Realistic Scene Generation with LiDAR Diffusion Models**.
+Please refer to our [paper](https://arxiv.org/abs/2404.00815), [project page](https://lidar-diffusion.github.io/), or [github](https://github.com/hancyran/LiDAR-Diffusion) for more details."""

data/config/semantic-kitti.yaml

@@ -0,0 +1,211 @@
+# This file is covered by the LICENSE file in the root of this project.
+labels: 
+  0 : "unlabeled"
+  1 : "outlier"
+  10: "car"
+  11: "bicycle"
+  13: "bus"
+  15: "motorcycle"
+  16: "on-rails"
+  18: "truck"
+  20: "other-vehicle"
+  30: "person"
+  31: "bicyclist"
+  32: "motorcyclist"
+  40: "road"
+  44: "parking"
+  48: "sidewalk"
+  49: "other-ground"
+  50: "building"
+  51: "fence"
+  52: "other-structure"
+  60: "lane-marking"
+  70: "vegetation"
+  71: "trunk"
+  72: "terrain"
+  80: "pole"
+  81: "traffic-sign"
+  99: "other-object"
+  252: "moving-car"
+  253: "moving-bicyclist"
+  254: "moving-person"
+  255: "moving-motorcyclist"
+  256: "moving-on-rails"
+  257: "moving-bus"
+  258: "moving-truck"
+  259: "moving-other-vehicle"
+color_map: # bgr
+  0 : [0, 0, 0]
+  1 : [0, 0, 255]
+  10: [245, 150, 100]
+  11: [245, 230, 100]
+  13: [250, 80, 100]
+  15: [150, 60, 30]
+  16: [255, 0, 0]
+  18: [180, 30, 80]
+  20: [255, 0, 0]
+  30: [30, 30, 255]
+  31: [200, 40, 255]
+  32: [90, 30, 150]
+  40: [255, 0, 255]
+  44: [255, 150, 255]
+  48: [75, 0, 75]
+  49: [75, 0, 175]
+  50: [0, 200, 255]
+  51: [50, 120, 255]
+  52: [0, 150, 255]
+  60: [170, 255, 150]
+  70: [0, 175, 0]
+  71: [0, 60, 135]
+  72: [80, 240, 150]
+  80: [150, 240, 255]
+  81: [0, 0, 255]
+  99: [255, 255, 50]
+  252: [245, 150, 100]
+  256: [255, 0, 0]
+  253: [200, 40, 255]
+  254: [30, 30, 255]
+  255: [90, 30, 150]
+  257: [250, 80, 100]
+  258: [180, 30, 80]
+  259: [255, 0, 0]
+content: # as a ratio with the total number of points
+  0: 0.018889854628292943
+  1: 0.0002937197336781505
+  10: 0.040818519255974316
+  11: 0.00016609538710764618
+  13: 2.7879693665067774e-05
+  15: 0.00039838616015114444
+  16: 0.0
+  18: 0.0020633612104619787
+  20: 0.0016218197275284021
+  30: 0.00017698551338515307
+  31: 1.1065903904919655e-08
+  32: 5.532951952459828e-09
+  40: 0.1987493871255525
+  44: 0.014717169549888214
+  48: 0.14392298360372
+  49: 0.0039048553037472045
+  50: 0.1326861944777486
+  51: 0.0723592229456223
+  52: 0.002395131480328884
+  60: 4.7084144280367186e-05
+  70: 0.26681502148037506
+  71: 0.006035012012626033
+  72: 0.07814222006271769
+  80: 0.002855498193863172
+  81: 0.0006155958086189918
+  99: 0.009923127583046915
+  252: 0.001789309418528068
+  253: 0.00012709999297008662
+  254: 0.00016059776092534436
+  255: 3.745553104802113e-05
+  256: 0.0
+  257: 0.00011351574470342043
+  258: 0.00010157861367183268
+  259: 4.3840131989471124e-05
+# classes that are indistinguishable from single scan or inconsistent in
+# ground truth are mapped to their closest equivalent
+learning_map:
+  0 : 0     # "unlabeled"
+  1 : 0     # "outlier" mapped to "unlabeled" --------------------------mapped
+  10: 1     # "car"
+  11: 2     # "bicycle"
+  13: 5     # "bus" mapped to "other-vehicle" --------------------------mapped
+  15: 3     # "motorcycle"
+  16: 5     # "on-rails" mapped to "other-vehicle" ---------------------mapped
+  18: 4     # "truck"
+  20: 5     # "other-vehicle"
+  30: 6     # "person"
+  31: 7     # "bicyclist"
+  32: 8     # "motorcyclist"
+  40: 9     # "road"
+  44: 10    # "parking"
+  48: 11    # "sidewalk"
+  49: 12    # "other-ground"
+  50: 13    # "building"
+  51: 14    # "fence"
+  52: 0     # "other-structure" mapped to "unlabeled" ------------------mapped
+  60: 9     # "lane-marking" to "road" ---------------------------------mapped
+  70: 15    # "vegetation"
+  71: 16    # "trunk"
+  72: 17    # "terrain"
+  80: 18    # "pole"
+  81: 19    # "traffic-sign"
+  99: 0     # "other-object" to "unlabeled" ----------------------------mapped
+  252: 1    # "moving-car" to "car" ------------------------------------mapped
+  253: 7    # "moving-bicyclist" to "bicyclist" ------------------------mapped
+  254: 6    # "moving-person" to "person" ------------------------------mapped
+  255: 8    # "moving-motorcyclist" to "motorcyclist" ------------------mapped
+  256: 5    # "moving-on-rails" mapped to "other-vehicle" --------------mapped
+  257: 5    # "moving-bus" mapped to "other-vehicle" -------------------mapped
+  258: 4    # "moving-truck" to "truck" --------------------------------mapped
+  259: 5    # "moving-other"-vehicle to "other-vehicle" ----------------mapped
+learning_map_inv: # inverse of previous map
+  0: 0      # "unlabeled", and others ignored
+  1: 10     # "car"
+  2: 11     # "bicycle"
+  3: 15     # "motorcycle"
+  4: 18     # "truck"
+  5: 20     # "other-vehicle"
+  6: 30     # "person"
+  7: 31     # "bicyclist"
+  8: 32     # "motorcyclist"
+  9: 40     # "road"
+  10: 44    # "parking"
+  11: 48    # "sidewalk"
+  12: 49    # "other-ground"
+  13: 50    # "building"
+  14: 51    # "fence"
+  15: 70    # "vegetation"
+  16: 71    # "trunk"
+  17: 72    # "terrain"
+  18: 80    # "pole"
+  19: 81    # "traffic-sign"
+learning_ignore: # Ignore classes
+  0: True      # "unlabeled", and others ignored
+  1: False     # "car"
+  2: False     # "bicycle"
+  3: False     # "motorcycle"
+  4: False     # "truck"
+  5: False     # "other-vehicle"
+  6: False     # "person"
+  7: False     # "bicyclist"
+  8: False     # "motorcyclist"
+  9: False     # "road"
+  10: False    # "parking"
+  11: False    # "sidewalk"
+  12: False    # "other-ground"
+  13: False    # "building"
+  14: False    # "fence"
+  15: False    # "vegetation"
+  16: False    # "trunk"
+  17: False    # "terrain"
+  18: False    # "pole"
+  19: False    # "traffic-sign"
+split: # sequence numbers
+  train:
+    - 0
+    - 1
+    - 2
+    - 3
+    - 4
+    - 5
+    - 6
+    - 7
+    - 9
+    - 10
+  valid:
+    - 8
+  test:
+    - 11
+    - 12
+    - 13
+    - 14
+    - 15
+    - 16
+    - 17
+    - 18
+    - 19
+    - 20
+    - 21

lidm/data/annotated_dataset.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+from typing import Optional, List, Dict, Union, Any
+import warnings
+
+from torch.utils.data import Dataset
+
+from .conditional_builder.objects_bbox import ObjectsBoundingBoxConditionalBuilder
+from .conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder
+
+
+class Annotated3DObjectsDataset(Dataset):
+    def __init__(self, min_objects_per_image: int,
+                 max_objects_per_image: int, no_tokens: int, num_beams: int, cats: List[str],
+                 cat_blacklist: Optional[List[str]] = None, **kwargs):
+        self.min_objects_per_image = min_objects_per_image
+        self.max_objects_per_image = max_objects_per_image
+        self.no_tokens = no_tokens
+        self.num_beams = num_beams
+
+        self.categories = [c for c in cats if c not in cat_blacklist] if cat_blacklist is not None else cats
+        self._conditional_builders = None
+
+    @property
+    def no_classes(self) -> int:
+        return len(self.categories)
+
+    @property
+    def conditional_builders(self) -> ObjectsCenterPointsConditionalBuilder:
+        # cannot set this up in init because no_classes is only known after loading data in init of superclass
+        if self._conditional_builders is None:
+            self._conditional_builders = {
+                'center': ObjectsCenterPointsConditionalBuilder(
+                    self.no_classes,
+                    self.max_objects_per_image,
+                    self.no_tokens,
+                    self.num_beams
+                ),
+                'bbox': ObjectsBoundingBoxConditionalBuilder(
+                    self.no_classes,
+                    self.max_objects_per_image,
+                    self.no_tokens,
+                    self.num_beams
+                )
+            }
+        return self._conditional_builders
+
+    def get_textual_label_for_category_id(self, category_id: int) -> str:
+        return self.categories[category_id]

lidm/data/base.py

@@ -0,0 +1,121 @@
+import pdb
+from abc import abstractmethod
+from functools import partial
+
+import PIL
+import numpy as np
+from PIL import Image
+
+import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset, IterableDataset
+
+from ..utils.aug_utils import get_lidar_transform, get_camera_transform, get_anno_transform
+
+
+class DatasetBase(Dataset):
+    def __init__(self, data_root, split, dataset_config, aug_config, return_pcd=False, condition_key=None,
+                 scale_factors=None, degradation=None, **kwargs):
+        self.data_root = data_root
+        self.split = split
+        self.data = []
+        self.aug_config = aug_config
+
+        self.img_size = dataset_config.size
+        self.fov = dataset_config.fov
+        self.depth_range = dataset_config.depth_range
+        self.filtered_map_cats = dataset_config.filtered_map_cats
+        self.depth_scale = dataset_config.depth_scale
+        self.log_scale = dataset_config.log_scale
+
+        if self.log_scale:
+            self.depth_thresh = (np.log2(1./255. + 1) / self.depth_scale) * 2. - 1 + 1e-6
+        else:
+            self.depth_thresh = (1./255. / self.depth_scale) * 2. - 1 + 1e-6
+        self.return_pcd = return_pcd
+
+        if degradation is not None and scale_factors is not None:
+            scaled_img_size = (int(self.img_size[0] / scale_factors[0]), int(self.img_size[1] / scale_factors[1]))
+            degradation_fn = {
+                "pil_nearest": PIL.Image.NEAREST,
+                "pil_bilinear": PIL.Image.BILINEAR,
+                "pil_bicubic": PIL.Image.BICUBIC,
+                "pil_box": PIL.Image.BOX,
+                "pil_hamming": PIL.Image.HAMMING,
+                "pil_lanczos": PIL.Image.LANCZOS,
+            }[degradation]
+            self.degradation_transform = partial(TF.resize, size=scaled_img_size, interpolation=degradation_fn)
+        else:
+            self.degradation_transform = None
+        self.condition_key = condition_key
+
+        self.lidar_transform = get_lidar_transform(aug_config, split)
+        self.anno_transform = get_anno_transform(aug_config, split) if condition_key in ['bbox', 'center'] else None
+        self.view_transform = get_camera_transform(aug_config, split) if condition_key in ['camera'] else None
+
+        self.prepare_data()
+
+    def prepare_data(self):
+        raise NotImplementedError
+
+    def process_scan(self, range_img):
+        range_img = np.where(range_img < 0, 0, range_img)
+
+        if self.log_scale:
+            # log scale
+            range_img = np.log2(range_img + 0.0001 + 1)
+
+        range_img = range_img / self.depth_scale
+        range_img = range_img * 2. - 1.
+
+        range_img = np.clip(range_img, -1, 1)
+        range_img = np.expand_dims(range_img, axis=0)
+
+        # mask
+        range_mask = np.ones_like(range_img)
+        range_mask[range_img < self.depth_thresh] = -1
+
+        return range_img, range_mask
+
+    @staticmethod
+    def load_lidar_sweep(*args, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    def load_semantic_map(*args, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    def load_camera(*args, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    def load_annotation(*args, **kwargs):
+        raise NotImplementedError
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        example = dict()
+        return example
+
+
+class Txt2ImgIterableBaseDataset(IterableDataset):
+    """
+    Define an interface to make the IterableDatasets for text2img data chainable
+    """
+    def __init__(self, num_records=0, valid_ids=None, size=256):
+        super().__init__()
+        self.num_records = num_records
+        self.valid_ids = valid_ids
+        self.sample_ids = valid_ids
+        self.size = size
+
+        print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
+
+    def __len__(self):
+        return self.num_records
+
+    @abstractmethod
+    def __iter__(self):
+        pass

lidm/data/conditional_builder/objects_bbox.py

@@ -0,0 +1,53 @@
+from itertools import cycle
+from typing import List, Tuple, Callable, Optional
+
+from PIL import Image as pil_image, ImageDraw as pil_img_draw, ImageFont
+from more_itertools.recipes import grouper
+from torch import LongTensor, Tensor
+
+from ..helper_types import BoundingBox, Annotation
+from .objects_center_points import ObjectsCenterPointsConditionalBuilder, convert_pil_to_tensor
+from .utils import COLOR_PALETTE, WHITE, GRAY_75, BLACK, additional_parameters_string, \
+    pad_list, get_plot_font_size, absolute_bbox
+
+
+class ObjectsBoundingBoxConditionalBuilder(ObjectsCenterPointsConditionalBuilder):
+    @property
+    def object_descriptor_length(self) -> int:
+        return 3  # 3/5: object_representation (1) + corners (2/4)
+
+    def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]:
+        object_tuples = [
+            (self.object_representation(ann), *self.token_pair_from_bbox(ann.bbox))
+            for ann in annotations
+        ]
+        object_tuples = pad_list(object_tuples, self.empty_tuple, self.no_max_objects)
+        return object_tuples
+
+    def inverse_build(self, conditional: LongTensor) -> Tuple[List[Tuple[int, BoundingBox]], Optional[BoundingBox]]:
+        conditional_list = conditional.tolist()
+        object_triples = grouper(conditional_list, 3)
+        assert conditional.shape[0] == self.embedding_dim
+        return [(object_triple[0], self.bbox_from_token_pair(object_triple[1], object_triple[2])) for object_triple in object_triples if object_triple[0] != self.none], None
+
+    def plot(self, conditional: LongTensor, label_for_category_no: Callable[[int], str], figure_size: Tuple[int, int],
+             line_width: int = 3, font_size: Optional[int] = None) -> Tensor:
+        plot = pil_image.new('RGB', figure_size, WHITE)
+        draw = pil_img_draw.Draw(plot)
+        # font = ImageFont.truetype(
+        #     "/usr/share/fonts/truetype/lato/Lato-Regular.ttf",
+        #     size=get_plot_font_size(font_size, figure_size)
+        # )
+        font = ImageFont.load_default()
+        width, height = plot.size
+        description, crop_coordinates = self.inverse_build(conditional)
+        for (representation, bbox), color in zip(description, cycle(COLOR_PALETTE)):
+            annotation = self.representation_to_annotation(representation)
+            # class_label = label_for_category_no(annotation.category_id) + ' ' + additional_parameters_string(annotation)
+            class_label = label_for_category_no(annotation.category_id)
+            bbox = absolute_bbox(bbox, width, height)
+            draw.rectangle(bbox, outline=color, width=line_width)
+            draw.text((bbox[0] + line_width, bbox[1] + line_width), class_label, anchor='la', fill=BLACK, font=font)
+        if crop_coordinates is not None:
+            draw.rectangle(absolute_bbox(crop_coordinates, width, height), outline=GRAY_75, width=line_width)
+        return convert_pil_to_tensor(plot) / 127.5 - 1.

lidm/data/conditional_builder/objects_center_points.py

@@ -0,0 +1,150 @@
+import math
+import random
+import warnings
+from itertools import cycle
+from typing import List, Optional, Tuple, Callable
+
+from PIL import Image as pil_image, ImageDraw as pil_img_draw, ImageFont
+from more_itertools.recipes import grouper
+from .utils import COLOR_PALETTE, WHITE, GRAY_75, BLACK, additional_parameters_string, pad_list, get_circle_size, \
+    get_plot_font_size, absolute_bbox
+from ..helper_types import BoundingBox, Annotation, Image
+from torch import LongTensor, Tensor
+from torchvision.transforms import PILToTensor
+
+
+pil_to_tensor = PILToTensor()
+
+
+def convert_pil_to_tensor(image: Image) -> Tensor:
+    with warnings.catch_warnings():
+        # to filter PyTorch UserWarning as described here: https://github.com/pytorch/vision/issues/2194
+        warnings.simplefilter("ignore")
+        return pil_to_tensor(image)
+
+
+class ObjectsCenterPointsConditionalBuilder:
+    def __init__(self, no_object_classes: int, no_max_objects: int, no_tokens: int, num_beams: int):
+        self.no_object_classes = no_object_classes
+        self.no_max_objects = no_max_objects
+        self.no_tokens = no_tokens
+        # self.no_sections = int(math.sqrt(self.no_tokens))
+        self.no_sections = (self.no_tokens // num_beams, num_beams)  # (width, height)
+
+    @property
+    def none(self) -> int:
+        return self.no_tokens - 1
+
+    @property
+    def object_descriptor_length(self) -> int:
+        return 2
+
+    @property
+    def empty_tuple(self) -> Tuple:
+        return (self.none,) * self.object_descriptor_length
+
+    @property
+    def embedding_dim(self) -> int:
+        return self.no_max_objects * self.object_descriptor_length
+
+    def tokenize_coordinates(self, x: float, y: float) -> int:
+        """
+        Express 2d coordinates with one number.
+        Example: assume self.no_tokens = 16, then no_sections = 4:
+        0  0  0  0
+        0  0  #  0
+        0  0  0  0
+        0  0  0  x
+        Then the # position corresponds to token 6, the x position to token 15.
+        @param x: float in [0, 1]
+        @param y: float in [0, 1]
+        @return: discrete tokenized coordinate
+        """
+        x_discrete = int(round(x * (self.no_sections[0] - 1)))
+        y_discrete = int(round(y * (self.no_sections[1] - 1)))
+        return y_discrete * self.no_sections[0] + x_discrete
+
+    def coordinates_from_token(self, token: int) -> (float, float):
+        x = token % self.no_sections[0]
+        y = token // self.no_sections[0]
+        return x / (self.no_sections[0] - 1), y / (self.no_sections[1] - 1)
+
+    def bbox_from_token_pair(self, token1: int, token2: int) -> BoundingBox:
+        x0, y0 = self.coordinates_from_token(token1)
+        x1, y1 = self.coordinates_from_token(token2)
+        # x2, y2 = self.coordinates_from_token(token3)
+        # x3, y3 = self.coordinates_from_token(token4)
+        return x0, y0, x1, y1
+
+    def token_pair_from_bbox(self, bbox: BoundingBox) -> Tuple:
+        # return self.tokenize_coordinates(bbox[0], bbox[1]), self.tokenize_coordinates(bbox[2], bbox[3]), self.tokenize_coordinates(bbox[4], bbox[5]), self.tokenize_coordinates(bbox[6], bbox[7])
+        return self.tokenize_coordinates(bbox[0], bbox[1]), self.tokenize_coordinates(bbox[4], bbox[5])
+
+    def inverse_build(self, conditional: LongTensor) \
+            -> Tuple[List[Tuple[int, Tuple[float, float]]], Optional[BoundingBox]]:
+        conditional_list = conditional.tolist()
+        table_of_content = grouper(conditional_list, self.object_descriptor_length)
+        assert conditional.shape[0] == self.embedding_dim
+        return [
+            (object_tuple[0], self.coordinates_from_token(object_tuple[1]))
+            for object_tuple in table_of_content if object_tuple[0] != self.none
+        ], None
+
+    def plot(self, conditional: LongTensor, label_for_category_no: Callable[[int], str], figure_size: Tuple[int, int],
+             line_width: int = 3, font_size: Optional[int] = None) -> Tensor:
+        plot = pil_image.new('RGB', figure_size, WHITE)
+        draw = pil_img_draw.Draw(plot)
+        circle_size = get_circle_size(figure_size)
+        # font = ImageFont.truetype('/usr/share/fonts/truetype/lato/Lato-Regular.ttf',
+        #                           size=get_plot_font_size(font_size, figure_size))
+        font = ImageFont.load_default()
+        width, height = plot.size
+        description, crop_coordinates = self.inverse_build(conditional)
+        for (representation, (x, y)), color in zip(description, cycle(COLOR_PALETTE)):
+            x_abs, y_abs = x * width, y * height
+            ann = self.representation_to_annotation(representation)
+            label = label_for_category_no(ann.category_id) + ' ' + additional_parameters_string(ann)
+            ellipse_bbox = [x_abs - circle_size, y_abs - circle_size, x_abs + circle_size, y_abs + circle_size]
+            draw.ellipse(ellipse_bbox, fill=color, width=0)
+            draw.text((x_abs, y_abs), label, anchor='md', fill=BLACK, font=font)
+        if crop_coordinates is not None:
+            draw.rectangle(absolute_bbox(crop_coordinates, width, height), outline=GRAY_75, width=line_width)
+        return convert_pil_to_tensor(plot) / 127.5 - 1.
+
+    def object_representation(self, annotation: Annotation) -> int:
+        return annotation.category_id
+
+    def representation_to_annotation(self, representation: int) -> Annotation:
+        category_id = representation % self.no_object_classes
+        # noinspection PyTypeChecker
+        return Annotation(
+            bbox=None,
+            category_id=category_id,
+        )
+
+    def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]:
+        object_tuples = [
+            (self.object_representation(a),
+             self.tokenize_coordinates(a.center[0], a.center[1]))
+            for a in annotations
+        ]
+        empty_tuple = (self.none, self.none)
+        object_tuples = pad_list(object_tuples, empty_tuple, self.no_max_objects)
+        return object_tuples
+
+    def build(self, annotations: List[Annotation]) \
+            -> LongTensor:
+        if len(annotations) == 0:
+            warnings.warn('Did not receive any annotations.')
+
+        random.shuffle(annotations)
+        if len(annotations) > self.no_max_objects:
+            warnings.warn('Received more annotations than allowed.')
+            annotations = annotations[:self.no_max_objects]
+
+        object_tuples = self._make_object_descriptors(annotations)
+        flattened = [token for tuple_ in object_tuples for token in tuple_]
+        assert len(flattened) == self.embedding_dim
+        assert all(0 <= value < self.no_tokens for value in flattened)
+
+        return LongTensor(flattened)

lidm/data/conditional_builder/utils.py

@@ -0,0 +1,188 @@
+import importlib
+from typing import List, Any, Tuple, Optional
+
+import numpy as np
+from ..helper_types import BoundingBox, Annotation
+
+# source: seaborn, color palette tab10
+COLOR_PALETTE = [(30, 118, 179), (255, 126, 13), (43, 159, 43), (213, 38, 39), (147, 102, 188),
+                 (139, 85, 74), (226, 118, 193), (126, 126, 126), (187, 188, 33), (22, 189, 206)]
+BLACK = (0, 0, 0)
+GRAY_75 = (63, 63, 63)
+GRAY_50 = (127, 127, 127)
+GRAY_25 = (191, 191, 191)
+WHITE = (255, 255, 255)
+FULL_CROP = (0., 0., 1., 1.)
+
+
+def corners_3d_to_2d(corners3d):
+    """
+    Args:
+        corners3d: (N, 8, 2)
+    Returns:
+        corners2d: (N, 4, 2)
+    """
+    # select pairs to reorganize
+    mask_0_3 = corners3d[:, 0:4, 0].argmax(1) // 2 != 0
+    mask_4_7 = corners3d[:, 4:8, 0].argmin(1) // 2 != 0
+
+    # reorganize corners in the order of (bottom-right, bottom-left)
+    corners3d[mask_0_3, 0:4] = corners3d[mask_0_3][:, [2, 3, 0, 1]]
+    # reorganize corners in the order of (top-left, top-right)
+    corners3d[mask_4_7, 4:8] = corners3d[mask_4_7][:, [2, 3, 0, 1]]
+
+    # calculate corners in order
+    bot_r = np.stack([corners3d[:, 0:2, 0].max(1), corners3d[:, 0:2, 1].min(1)], axis=-1)
+    bot_l = np.stack([corners3d[:, 2:4, 0].min(1), corners3d[:, 2:4, 1].min(1)], axis=-1)
+    top_l = np.stack([corners3d[:, 4:6, 0].min(1), corners3d[:, 4:6, 1].max(1)], axis=-1)
+    top_r = np.stack([corners3d[:, 6:8, 0].max(1), corners3d[:, 6:8, 1].max(1)], axis=-1)
+
+    return np.stack([bot_r, bot_l, top_l, top_r], axis=1)
+
+
+def rotate_points_along_z(points, angle):
+    """
+    Args:
+        points: (N, 3 + C)
+        angle: angle along z-axis, angle increases x ==> y
+    Returns:
+
+    """
+    cosa = np.cos(angle)
+    sina = np.sin(angle)
+    zeros = np.zeros(points.shape[0])
+    ones = np.ones(points.shape[0])
+    rot_matrix = np.stack((
+        cosa,  sina, zeros,
+        -sina, cosa, zeros,
+        zeros, zeros, ones)).reshape((-1, 3, 3))
+    points_rot = np.matmul(points[:, :, 0:3], rot_matrix)
+    points_rot = np.concatenate((points_rot, points[:, :, 3:]), axis=-1)
+    return points_rot
+
+
+def boxes_to_corners_3d(boxes3d):
+    """
+        7 -------- 4
+       /|         /|
+      6 -------- 5 .
+      | |        | |
+      . 3 -------- 0
+      |/         |/
+      2 -------- 1
+    Args:
+        boxes3d:  (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center
+
+    Returns:
+        corners3d: (N, 8, 3)
+    """
+    template = np.array(
+        [[1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1],
+        [1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1]],
+    ) / 2
+
+    # corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
+    corners3d = np.tile(boxes3d[:, None, 3:6], (1, 8, 1)) * template[None, :, :]
+    corners3d = rotate_points_along_z(corners3d.reshape((-1, 8, 3)), boxes3d[:, 6]).reshape((-1, 8, 3))
+    corners3d += boxes3d[:, None, 0:3]
+
+    return corners3d
+
+
+def intersection_area(rectangle1: BoundingBox, rectangle2: BoundingBox) -> float:
+    """
+    Give intersection area of two rectangles.
+    @param rectangle1: (x0, y0, w, h) of first rectangle
+    @param rectangle2: (x0, y0, w, h) of second rectangle
+    """
+    rectangle1 = rectangle1[0], rectangle1[1], rectangle1[0] + rectangle1[2], rectangle1[1] + rectangle1[3]
+    rectangle2 = rectangle2[0], rectangle2[1], rectangle2[0] + rectangle2[2], rectangle2[1] + rectangle2[3]
+    x_overlap = max(0., min(rectangle1[2], rectangle2[2]) - max(rectangle1[0], rectangle2[0]))
+    y_overlap = max(0., min(rectangle1[3], rectangle2[3]) - max(rectangle1[1], rectangle2[1]))
+    return x_overlap * y_overlap
+
+
+def horizontally_flip_bbox(bbox: BoundingBox) -> BoundingBox:
+    return 1 - (bbox[0] + bbox[2]), bbox[1], bbox[2], bbox[3]
+
+
+def absolute_bbox(relative_bbox: BoundingBox, width: int, height: int) -> Tuple[int, int, int, int]:
+    bbox = relative_bbox
+    # bbox = bbox[0] * width, bbox[1] * height, (bbox[0] + bbox[2]) * width, (bbox[1] + bbox[3]) * height
+    bbox = bbox[0] * width, bbox[1] * height, bbox[2] * width, bbox[3] * height
+    # return int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
+    x1, x2 = min(int(bbox[2]), int(bbox[0])), max(int(bbox[2]), int(bbox[0]))
+    y1, y2 = min(int(bbox[3]), int(bbox[1])), max(int(bbox[3]), int(bbox[1]))
+    if x1 == x2:
+        x2 += 1
+    if y1 == y2:
+        y2 += 1
+    return x1, y1, x2, y2
+
+
+def pad_list(list_: List, pad_element: Any, pad_to_length: int) -> List:
+    return list_ + [pad_element for _ in range(pad_to_length - len(list_))]
+
+
+def rescale_annotations(annotations: List[Annotation], crop_coordinates: BoundingBox, flip: bool) -> \
+        List[Annotation]:
+    def clamp(x: float):
+        return max(min(x, 1.), 0.)
+
+    def rescale_bbox(bbox: BoundingBox) -> BoundingBox:
+        x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
+        y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
+        w = min(bbox[2] / crop_coordinates[2], 1 - x0)
+        h = min(bbox[3] / crop_coordinates[3], 1 - y0)
+        if flip:
+            x0 = 1 - (x0 + w)
+        return x0, y0, w, h
+
+    return [a._replace(bbox=rescale_bbox(a.bbox)) for a in annotations]
+
+
+def filter_annotations(annotations: List[Annotation], crop_coordinates: BoundingBox) -> List:
+    return [a for a in annotations if intersection_area(a.bbox, crop_coordinates) > 0.0]
+
+
+def additional_parameters_string(annotation: Annotation, short: bool = True) -> str:
+    sl = slice(1) if short else slice(None)
+    string = ''
+    if not (annotation.is_group_of or annotation.is_occluded or annotation.is_depiction or annotation.is_inside):
+        return string
+    if annotation.is_group_of:
+        string += 'group'[sl] + ','
+    if annotation.is_occluded:
+        string += 'occluded'[sl] + ','
+    if annotation.is_depiction:
+        string += 'depiction'[sl] + ','
+    if annotation.is_inside:
+        string += 'inside'[sl]
+    return '(' + string.strip(",") + ')'
+
+
+def get_plot_font_size(font_size: Optional[int], figure_size: Tuple[int, int]) -> int:
+    if font_size is None:
+        font_size = 10
+        if max(figure_size) >= 256:
+            font_size = 12
+        if max(figure_size) >= 512:
+            font_size = 15
+    return font_size
+
+
+def get_circle_size(figure_size: Tuple[int, int]) -> int:
+    circle_size = 2
+    if max(figure_size) >= 256:
+        circle_size = 3
+    if max(figure_size) >= 512:
+        circle_size = 4
+    return circle_size
+
+
+def load_object_from_string(object_string: str) -> Any:
+    """
+    Source: https://stackoverflow.com/a/10773699
+    """
+    module_name, class_name = object_string.rsplit(".", 1)
+    return getattr(importlib.import_module(module_name), class_name)

lidm/data/helper_types.py

@@ -0,0 +1,20 @@
+from typing import Tuple, Optional, NamedTuple, Union, List
+from PIL.Image import Image as pil_image
+from torch import Tensor
+
+try:
+  from typing import Literal
+except ImportError:
+  from typing_extensions import Literal
+
+Image = Union[Tensor, pil_image]
+# BoundingBox = Tuple[float, float, float, float]  # x0, y0, w, h | x0, y0, x1, y1
+# BoundingBox3D = Tuple[float, float, float, float, float, float]  # x0, y0, z0, l, w, h
+BoundingBox = Tuple[float, float, float, float]  # corner coordinates (x,y) in the order of bottom-right -> bottom-left -> top-left -> top-right
+Center = Tuple[float, float]
+
+
+class Annotation(NamedTuple):
+    category_id: int
+    bbox: Optional[BoundingBox] = None
+    center: Optional[Center] = None

lidm/data/kitti.py

@@ -0,0 +1,345 @@
+import glob
+import os
+import pickle
+import numpy as np
+import yaml
+from PIL import Image
+import xml.etree.ElementTree as ET
+
+from lidm.data.base import DatasetBase
+from .annotated_dataset import Annotated3DObjectsDataset
+from .conditional_builder.utils import corners_3d_to_2d
+from .helper_types import Annotation
+from ..utils.lidar_utils import pcd2range, pcd2coord2d, range2pcd
+
+# TODO add annotation categories and semantic categories
+CATEGORIES = ['ignore', 'car', 'bicycle', 'motorcycle', 'truck', 'other-vehicle', 'person', 'bicyclist', 'motorcyclist',
+              'road', 'parking', 'sidewalk', 'other-ground', 'building', 'fence', 'vegetation', 'trunk', 'terrain',
+              'pole', 'traffic-sign']
+CATE2LABEL = {k: v for v, k in enumerate(CATEGORIES)}  # 0: invalid, 1~10: categories
+LABEL2RGB = np.array([(0, 0, 0), (0, 0, 142), (119, 11, 32), (0, 0, 230), (0, 0, 70), (0, 0, 90), (220, 20, 60),
+                      (255, 0, 0), (0, 0, 110), (128, 64, 128), (250, 170, 160), (244, 35, 232), (230, 150, 140),
+                      (70, 70, 70), (190, 153, 153), (107, 142, 35), (0, 80, 100), (230, 150, 140), (153, 153, 153),
+                      (220, 220, 0)])
+CAMERAS = ['CAM_FRONT']
+BBOX_CATS = ['car', 'people', 'cycle']
+BBOX_CAT2LABEL = {'car': 0, 'truck': 0, 'bus': 0, 'caravan': 0, 'person': 1, 'rider': 2, 'motorcycle': 2, 'bicycle': 2}
+
+# train + test
+SEM_KITTI_TRAIN_SET = ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']
+KITTI_TRAIN_SET = SEM_KITTI_TRAIN_SET + ['11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']
+KITTI360_TRAIN_SET = ['00', '02', '04', '05', '06', '07', '09', '10'] + ['08']  # partial test data at '02' sequence
+CAM_KITTI360_TRAIN_SET = ['00', '04', '05', '06', '07', '08', '09', '10']  # cam mismatch lidar in '02'
+
+# validation
+SEM_KITTI_VAL_SET = KITTI_VAL_SET = ['08']
+CAM_KITTI360_VAL_SET = KITTI360_VAL_SET = ['03']
+
+
+class KITTIBase(DatasetBase):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dataset_name = 'kitti'
+        self.num_sem_cats = kwargs['dataset_config'].num_sem_cats + 1
+
+    @staticmethod
+    def load_lidar_sweep(path):
+        scan = np.fromfile(path, dtype=np.float32)
+        scan = scan.reshape((-1, 4))
+        points = scan[:, 0:3]  # get xyz
+        return points
+
+    def load_semantic_map(self, path, pcd):
+        raise NotImplementedError
+
+    def load_camera(self, path):
+        raise NotImplementedError
+
+    def __getitem__(self, idx):
+        example = dict()
+        data_path = self.data[idx]
+        # lidar point cloud
+        sweep = self.load_lidar_sweep(data_path)
+
+        if self.lidar_transform:
+            sweep, _ = self.lidar_transform(sweep, None)
+
+        if self.condition_key == 'segmentation':
+            # semantic maps
+            proj_range, sem_map = self.load_semantic_map(data_path, sweep)
+            example[self.condition_key] = sem_map
+        else:
+            proj_range, _ = pcd2range(sweep, self.img_size, self.fov, self.depth_range)
+        proj_range, proj_mask = self.process_scan(proj_range)
+        example['image'], example['mask'] = proj_range, proj_mask
+        if self.return_pcd:
+            reproj_sweep, _, _ = range2pcd(proj_range[0] * .5 + .5, self.fov, self.depth_range, self.depth_scale, self.log_scale)
+            example['raw'] = sweep
+            example['reproj'] = reproj_sweep.astype(np.float32)
+
+        # image degradation
+        if self.degradation_transform:
+            degraded_proj_range = self.degradation_transform(proj_range)
+            example['degraded_image'] = degraded_proj_range
+
+        # cameras
+        if self.condition_key == 'camera':
+            cameras = self.load_camera(data_path)
+            example[self.condition_key] = cameras
+
+        return example
+
+
+class SemanticKITTIBase(KITTIBase):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        assert self.condition_key in ['segmentation']  # for segmentation input only
+        self.label2rgb = LABEL2RGB
+
+    def prepare_data(self):
+        # read data paths from KITTI
+        for seq_id in eval('SEM_KITTI_%s_SET' % self.split.upper()):
+            self.data.extend(glob.glob(os.path.join(
+                self.data_root, f'dataset/sequences/{seq_id}/velodyne/*.bin')))
+        # read label mapping
+        data_config = yaml.safe_load(open('./data/config/semantic-kitti.yaml', 'r'))
+        remap_dict = data_config["learning_map"]
+        max_key = max(remap_dict.keys())
+        self.learning_map = np.zeros((max_key + 100), dtype=np.int32)
+        self.learning_map[list(remap_dict.keys())] = list(remap_dict.values())
+
+    def load_semantic_map(self, path, pcd):
+        label_path = path.replace('velodyne', 'labels').replace('.bin', '.label')
+        labels = np.fromfile(label_path, dtype=np.uint32)
+        labels = labels.reshape((-1))
+        labels = labels & 0xFFFF  # semantic label in lower half
+        labels = self.learning_map[labels]
+
+        proj_range, sem_map = pcd2range(pcd, self.img_size, self.fov, self.depth_range, labels=labels)
+        # sem_map = np.expand_dims(sem_map, axis=0).astype(np.int64)
+        sem_map = sem_map.astype(np.int64)
+        if self.filtered_map_cats is not None:
+            sem_map[np.isin(sem_map, self.filtered_map_cats)] = 0  # set filtered category as noise
+        onehot = np.eye(self.num_sem_cats, dtype=np.float32)[sem_map].transpose(2, 0, 1)
+        return proj_range, onehot
+
+
+class SemanticKITTITrain(SemanticKITTIBase):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/SemanticKITTI', split='train', **kwargs)
+
+
+class SemanticKITTIValidation(SemanticKITTIBase):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/SemanticKITTI', split='val', **kwargs)
+
+
+class KITTI360Base(KITTIBase):
+    def __init__(self, split_per_view=None, **kwargs):
+        super().__init__(**kwargs)
+        self.split_per_view = split_per_view
+        if self.condition_key == 'camera':
+            assert self.split_per_view is not None, 'For camera-to-lidar, need to specify split_per_view'
+
+    def prepare_data(self):
+        # read data paths
+        self.data = []
+        if self.condition_key == 'camera':
+            seq_list = eval('CAM_KITTI360_%s_SET' % self.split.upper())
+        else:
+            seq_list = eval('KITTI360_%s_SET' % self.split.upper())
+        for seq_id in seq_list:
+            self.data.extend(glob.glob(os.path.join(
+                self.data_root, f'data_3d_raw/2013_05_28_drive_00{seq_id}_sync/velodyne_points/data/*.bin')))
+
+    def random_drop_camera(self, camera_list):
+        if np.random.rand() < self.aug_config['camera_drop'] and self.split == 'train':
+            camera_list = [np.zeros_like(c) if i != len(camera_list) // 2 else c for i, c in enumerate(camera_list)]  # keep the middle view only
+        return camera_list
+
+    def load_camera(self, path):
+        camera_path = path.replace('data_3d_raw', 'data_2d_camera').replace('velodyne_points/data', 'image_00/data_rect').replace('.bin', '.png')
+        camera = np.array(Image.open(camera_path)).astype(np.float32) / 255.
+        camera = camera.transpose(2, 0, 1)
+        if self.view_transform:
+            camera = self.view_transform(camera)
+        camera_list = np.split(camera, self.split_per_view, axis=2)  # split into n chunks as different views
+        camera_list = self.random_drop_camera(camera_list)
+        return camera_list
+
+
+class KITTI360Train(KITTI360Base):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/KITTI-360', split='train', **kwargs)
+
+
+class KITTI360Validation(KITTI360Base):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/KITTI-360', split='val', **kwargs)
+
+
+class AnnotatedKITTI360Base(Annotated3DObjectsDataset, KITTI360Base):
+    def __init__(self, **kwargs):
+        self.id_bbox_dict = dict()
+        self.id_label_dict = dict()
+
+        Annotated3DObjectsDataset.__init__(self, **kwargs)
+        KITTI360Base.__init__(self, **kwargs)
+        assert self.condition_key in ['center', 'bbox']  # for annotated images only
+
+    @staticmethod
+    def parseOpencvMatrix(node):
+        rows = int(node.find('rows').text)
+        cols = int(node.find('cols').text)
+        data = node.find('data').text.split(' ')
+
+        mat = []
+        for d in data:
+            d = d.replace('\n', '')
+            if len(d) < 1:
+                continue
+            mat.append(float(d))
+        mat = np.reshape(mat, [rows, cols])
+        return mat
+
+    def parseVertices(self, child):
+        transform = self.parseOpencvMatrix(child.find('transform'))
+        R = transform[:3, :3]
+        T = transform[:3, 3]
+        vertices = self.parseOpencvMatrix(child.find('vertices'))
+        vertices = np.matmul(R, vertices.transpose()).transpose() + T
+        return vertices
+
+    def parse_bbox_xml(self, path):
+        tree = ET.parse(path)
+        root = tree.getroot()
+
+        bbox_dict = dict()
+        label_dict = dict()
+        for child in root:
+            if child.find('transform') is None:
+                continue
+
+            label_name = child.find('label').text
+            if label_name not in BBOX_CAT2LABEL:
+                continue
+
+            label = BBOX_CAT2LABEL[label_name]
+            timestamp = int(child.find('timestamp').text)
+            # verts = self.parseVertices(child)
+            verts = self.parseOpencvMatrix(child.find('vertices'))[:8]
+            if timestamp in bbox_dict:
+                bbox_dict[timestamp].append(verts)
+                label_dict[timestamp].append(label)
+            else:
+                bbox_dict[timestamp] = [verts]
+                label_dict[timestamp] = [label]
+        return bbox_dict, label_dict
+
+    def prepare_data(self):
+        KITTI360Base.prepare_data(self)
+
+        self.data = [p for p in self.data if '2013_05_28_drive_0008_sync' not in p]  # remove unlabeled sequence 08
+        seq_list = eval('KITTI360_%s_SET' % self.split.upper())
+        for seq_id in seq_list:
+            if seq_id != '08':
+                xml_path = os.path.join(self.data_root, f'data_3d_bboxes/train/2013_05_28_drive_00{seq_id}_sync.xml')
+                bbox_dict, label_dict = self.parse_bbox_xml(xml_path)
+                self.id_bbox_dict[seq_id] = bbox_dict
+                self.id_label_dict[seq_id] = label_dict
+
+    def load_annotation(self, path):
+        seq_id = path.split('/')[-4].split('_')[-2][-2:]
+        timestamp = int(path.split('/')[-1].replace('.bin', ''))
+        verts_list = self.id_bbox_dict[seq_id][timestamp]
+        label_list = self.id_label_dict[seq_id][timestamp]
+
+        if self.condition_key == 'bbox':
+            points = np.stack(verts_list)
+        elif self.condition_key == 'center':
+            points = (verts_list[0] + verts_list[6]) / 2.
+        else:
+            raise NotImplementedError
+        labels = np.array([label_list])
+        if self.anno_transform:
+            points, labels = self.anno_transform(points, labels)
+        return points, labels
+
+    def __getitem__(self, idx):
+        example = dict()
+        data_path = self.data[idx]
+
+        # lidar point cloud
+        sweep = self.load_lidar_sweep(data_path)
+
+        # annotations
+        bbox_points, bbox_labels = self.load_annotation(data_path)
+
+        if self.lidar_transform:
+            sweep, bbox_points = self.lidar_transform(sweep, bbox_points)
+
+        # point cloud -> range
+        proj_range, _ = pcd2range(sweep, self.img_size, self.fov, self.depth_range)
+        proj_range, proj_mask = self.process_scan(proj_range)
+        example['image'], example['mask'] = proj_range, proj_mask
+        if self.return_pcd:
+            example['reproj'] = sweep
+
+        # annotation -> range
+        # NOTE: do not need to transform bbox points along with lidar, since their coordinates are based on range-image space instead of 3D space
+        proj_bbox_points, proj_bbox_labels = pcd2coord2d(bbox_points, self.fov, self.depth_range, labels=bbox_labels)
+        builder = self.conditional_builders[self.condition_key]
+        if self.condition_key == 'bbox':
+            proj_bbox_points = corners_3d_to_2d(proj_bbox_points)
+            annotations = [Annotation(bbox=bbox.flatten(), category_id=label) for bbox, label in
+                           zip(proj_bbox_points, proj_bbox_labels)]
+        else:
+            annotations = [Annotation(center=center, category_id=label) for center, label in
+                           zip(proj_bbox_points, proj_bbox_labels)]
+        example[self.condition_key] = builder.build(annotations)
+
+        return example
+
+
+class AnnotatedKITTI360Train(AnnotatedKITTI360Base):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/KITTI-360', split='train', cats=BBOX_CATS, **kwargs)
+
+
+class AnnotatedKITTI360Validation(AnnotatedKITTI360Base):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset/KITTI-360', split='train', cats=BBOX_CATS, **kwargs)
+
+
+class KITTIImageBase(KITTIBase):
+    """
+    Range ImageSet only combining KITTI-360 and SemanticKITTI
+
+    #Samples (Training): 98014, #Samples (Val): 3511
+
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        assert self.condition_key in [None, 'image']  # for image input only
+
+    def prepare_data(self):
+        # read data paths from KITTI-360
+        self.data = []
+        for seq_id in eval('KITTI360_%s_SET' % self.split.upper()):
+            self.data.extend(glob.glob(os.path.join(
+                self.data_root, f'KITTI-360/data_3d_raw/2013_05_28_drive_00{seq_id}_sync/velodyne_points/data/*.bin')))
+
+        # read data paths from KITTI
+        for seq_id in eval('KITTI_%s_SET' % self.split.upper()):
+            self.data.extend(glob.glob(os.path.join(
+                self.data_root, f'SemanticKITTI/dataset/sequences/{seq_id}/velodyne/*.bin')))
+
+
+class KITTIImageTrain(KITTIImageBase):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset', split='train', **kwargs)
+
+
+class KITTIImageValidation(KITTIImageBase):
+    def __init__(self, **kwargs):
+        super().__init__(data_root='./dataset', split='val', **kwargs)

lidm/eval/README.md

@@ -0,0 +1,95 @@
+# 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

@@ -0,0 +1,62 @@
+"""
+@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

@@ -0,0 +1,9 @@
+#!/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

@@ -0,0 +1,138 @@
+"""
+@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

@@ -0,0 +1,191 @@
+"""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

@@ -0,0 +1,458 @@
+"""
+@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

@@ -0,0 +1,141 @@
+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

@@ -0,0 +1,372 @@
+#!/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

@@ -0,0 +1,179 @@
+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

@@ -0,0 +1,79 @@
+#!/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

@@ -0,0 +1,90 @@
+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

@@ -0,0 +1,182 @@
+
+#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

@@ -0,0 +1,33 @@
+#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

@@ -0,0 +1,84 @@
+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

@@ -0,0 +1,14 @@
+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

@@ -0,0 +1,196 @@
+
+#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

@@ -0,0 +1,33 @@
+#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

@@ -0,0 +1,76 @@
+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

@@ -0,0 +1,14 @@
+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

@@ -0,0 +1,31 @@
+// 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

@@ -0,0 +1,316 @@
+// 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

@@ -0,0 +1,112 @@
+# 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

@@ -0,0 +1,14 @@
+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
+    })