ocnn 2.2.5__tar.gz → 2.2.7__tar.gz

{ocnn-2.2.5 → ocnn-2.2.7}/LICENSE

@@ -1,21 +1,21 @@
-MIT License
-
-Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in
-all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
+MIT License
+
+Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

{ocnn-2.2.5 → ocnn-2.2.7}/MANIFEST.in

@@ -1 +1 @@
-recursive-exclude test *
+recursive-exclude test *

ocnn-2.2.7/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: ocnn
+Version: 2.2.7
+Summary: Octree-based Sparse Convolutional Neural Networks
+Home-page: https://github.com/octree-nn/ocnn-pytorch
+Author: Peng-Shuai Wang
+Author-email: wangps@hotmail.com
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: numpy
+Requires-Dist: packaging
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# O-CNN
+
+**[Documentation](https://ocnn-pytorch.readthedocs.io)**
+
+[![Documentation Status](https://readthedocs.org/projects/ocnn-pytorch/badge/?version=latest)](https://ocnn-pytorch.readthedocs.io/en/latest/?badge=latest)
+[![Downloads](https://static.pepy.tech/badge/ocnn)](https://pepy.tech/project/ocnn)
+[![Downloads](https://static.pepy.tech/badge/ocnn/month)](https://pepy.tech/project/ocnn)
+[![PyPI](https://img.shields.io/pypi/v/ocnn)](https://pypi.org/project/ocnn/)
+
+This repository contains the **pure PyTorch**-based implementation of
+[O-CNN](https://wang-ps.github.io/O-CNN.html). The code has been tested with
+`Pytorch>=1.6.0`, and `Pytorch>=1.9.0` is preferred. The *original*
+implementation of O-CNN is based on C++ and CUDA and can be found
+[here](https://github.com/Microsoft/O-CNN), which has received
+[![stars - O-CNN](https://img.shields.io/github/stars/microsoft/O-CNN?style=social)](https://github.com/microsoft/O-CNN) and
+[![forks - O-CNN](https://img.shields.io/github/forks/microsoft/O-CNN?style=social)](https://github.com/microsoft/O-CNN).
+
+
+O-CNN is an octree-based 3D convolutional neural network framework for 3D data.
+O-CNN constrains the CNN storage and computation into non-empty sparse voxels
+for efficiency and uses the `octree` data structure to organize and index these
+sparse voxels. Currently, this type of 3D convolution is known as Sparse
+Convolution in the research community.
+
+
+The concept of Sparse Convolution in O-CNN is the same with
+[SparseConvNet](https://openaccess.thecvf.com/content_cvpr_2018/papers/Graham_3D_Semantic_Segmentation_CVPR_2018_paper.pdf),
+[MinkowskiNet](https://github.com/NVIDIA/MinkowskiEngine), and
+[SpConv](https://github.com/traveller59/spconv).
+The key difference is that our O-CNN uses `octrees` to index the sparse voxels,
+while these works use `Hash Tables`. However, I believe that `octrees` may be
+the right choice for Sparse Convolution. With `octrees`, I can implement the
+Sparse Convolution with pure PyTorch. More importantly, with `octrees`, I can
+also build efficient transformers for 3D data --
+[OctFormer](https://github.com/octree-nn/octformer), which is extremely hard
+with `Hash Tables`.
+
+
+Our O-CNN is published in SIGGRAPH 2017, SparseConvNet is published in CVPR
+2018, and MinkowskiNet is published in CVPR 2019. Actually, our O-CNN was
+submitted to SIGGRAPH in the end of 2016 and was officially accepted in March,
+2017. <!-- The camera-ready version of our O-CNN was submitted to SIGGRAPH in April, 2018. -->
+We just did not post our paper on Arxiv during the review process of SIGGRAPH.
+Therefore, **the idea of constraining CNN computation into sparse non-emtpry
+voxels, i.e. Sparse Convolution,  is first proposed by our O-CNN**.
+
+<!--
+Developed in collaboration with authors from [PointCNN](https://arxiv.org/abs/1801.07791),
+[Dr. Yangyan Li](https://yangyan.li/) and [Prof. Baoquan Chen](https://baoquanchen.info/),
+-->
+This library supports point cloud processing from the ground up.
+The library provides essential components for converting raw point clouds into
+octrees to perform convolution operations. Of course, it also supports other 3D
+data formats, such as meshes and volumetric grids, which can be converted into
+octrees to leverage the library's capabilities.
+
+
+## Key benefits of ocnn-pytorch
+
+- **Simplicity**. The ocnn-pytorch is based on pure PyTorch, it is portable and
+  can be installed with a simple command:`pip install ocnn`. Other sparse
+  convolution frameworks heavily rely on C++ and CUDA, and it is complicated to
+  configure the compiling environment.
+
+- **Efficiency**. The ocnn-pytorch is very efficient compared with other sparse
+  convolution frameworks.  It only takes 18 hours to train the network on
+  ScanNet for 600 epochs with 4 V100 GPUs. For reference, under the same
+  training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
+  takes 30 hours.
+
+## Citation
+
+  ```bibtex
+  @article {Wang-2017-ocnn,
+    title    = {{O-CNN}: Octree-based Convolutional Neural Networksfor {3D} Shape Analysis},
+    author   = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
+    journal  = {ACM Transactions on Graphics (SIGGRAPH)},
+    volume   = {36},
+    number   = {4},
+    year     = {2017},
+  }
+  ```

ocnn-2.2.7/README.md

@@ -0,0 +1,82 @@
+# O-CNN
+
+**[Documentation](https://ocnn-pytorch.readthedocs.io)**
+
+[![Documentation Status](https://readthedocs.org/projects/ocnn-pytorch/badge/?version=latest)](https://ocnn-pytorch.readthedocs.io/en/latest/?badge=latest)
+[![Downloads](https://static.pepy.tech/badge/ocnn)](https://pepy.tech/project/ocnn)
+[![Downloads](https://static.pepy.tech/badge/ocnn/month)](https://pepy.tech/project/ocnn)
+[![PyPI](https://img.shields.io/pypi/v/ocnn)](https://pypi.org/project/ocnn/)
+
+This repository contains the **pure PyTorch**-based implementation of
+[O-CNN](https://wang-ps.github.io/O-CNN.html). The code has been tested with
+`Pytorch>=1.6.0`, and `Pytorch>=1.9.0` is preferred. The *original*
+implementation of O-CNN is based on C++ and CUDA and can be found
+[here](https://github.com/Microsoft/O-CNN), which has received
+[![stars - O-CNN](https://img.shields.io/github/stars/microsoft/O-CNN?style=social)](https://github.com/microsoft/O-CNN) and
+[![forks - O-CNN](https://img.shields.io/github/forks/microsoft/O-CNN?style=social)](https://github.com/microsoft/O-CNN).
+
+
+O-CNN is an octree-based 3D convolutional neural network framework for 3D data.
+O-CNN constrains the CNN storage and computation into non-empty sparse voxels
+for efficiency and uses the `octree` data structure to organize and index these
+sparse voxels. Currently, this type of 3D convolution is known as Sparse
+Convolution in the research community.
+
+
+The concept of Sparse Convolution in O-CNN is the same with
+[SparseConvNet](https://openaccess.thecvf.com/content_cvpr_2018/papers/Graham_3D_Semantic_Segmentation_CVPR_2018_paper.pdf),
+[MinkowskiNet](https://github.com/NVIDIA/MinkowskiEngine), and
+[SpConv](https://github.com/traveller59/spconv).
+The key difference is that our O-CNN uses `octrees` to index the sparse voxels,
+while these works use `Hash Tables`. However, I believe that `octrees` may be
+the right choice for Sparse Convolution. With `octrees`, I can implement the
+Sparse Convolution with pure PyTorch. More importantly, with `octrees`, I can
+also build efficient transformers for 3D data --
+[OctFormer](https://github.com/octree-nn/octformer), which is extremely hard
+with `Hash Tables`.
+
+
+Our O-CNN is published in SIGGRAPH 2017, SparseConvNet is published in CVPR
+2018, and MinkowskiNet is published in CVPR 2019. Actually, our O-CNN was
+submitted to SIGGRAPH in the end of 2016 and was officially accepted in March,
+2017. <!-- The camera-ready version of our O-CNN was submitted to SIGGRAPH in April, 2018. -->
+We just did not post our paper on Arxiv during the review process of SIGGRAPH.
+Therefore, **the idea of constraining CNN computation into sparse non-emtpry
+voxels, i.e. Sparse Convolution,  is first proposed by our O-CNN**.
+
+<!--
+Developed in collaboration with authors from [PointCNN](https://arxiv.org/abs/1801.07791),
+[Dr. Yangyan Li](https://yangyan.li/) and [Prof. Baoquan Chen](https://baoquanchen.info/),
+-->
+This library supports point cloud processing from the ground up.
+The library provides essential components for converting raw point clouds into
+octrees to perform convolution operations. Of course, it also supports other 3D
+data formats, such as meshes and volumetric grids, which can be converted into
+octrees to leverage the library's capabilities.
+
+
+## Key benefits of ocnn-pytorch
+
+- **Simplicity**. The ocnn-pytorch is based on pure PyTorch, it is portable and
+  can be installed with a simple command:`pip install ocnn`. Other sparse
+  convolution frameworks heavily rely on C++ and CUDA, and it is complicated to
+  configure the compiling environment.
+
+- **Efficiency**. The ocnn-pytorch is very efficient compared with other sparse
+  convolution frameworks.  It only takes 18 hours to train the network on
+  ScanNet for 600 epochs with 4 V100 GPUs. For reference, under the same
+  training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
+  takes 30 hours.
+
+## Citation
+
+  ```bibtex
+  @article {Wang-2017-ocnn,
+    title    = {{O-CNN}: Octree-based Convolutional Neural Networksfor {3D} Shape Analysis},
+    author   = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
+    journal  = {ACM Transactions on Graphics (SIGGRAPH)},
+    volume   = {36},
+    number   = {4},
+    year     = {2017},
+  }
+  ```

{ocnn-2.2.5 → ocnn-2.2.7}/ocnn/__init__.py

@@ -1,24 +1,24 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-from . import octree
-from . import nn
-from . import modules
-from . import models
-from . import dataset
-from . import utils
-
-__version__ = '2.2.5'
-
-__all__ = [
-    'octree',
-    'nn',
-    'modules',
-    'models',
-    'dataset',
-    'utils'
-]
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+from . import octree
+from . import nn
+from . import modules
+from . import models
+from . import dataset
+from . import utils
+
+__version__ = '2.2.7'
+
+__all__ = [
+    'octree',
+    'nn',
+    'modules',
+    'models',
+    'dataset',
+    'utils'
+]

{ocnn-2.2.5 → ocnn-2.2.7}/ocnn/dataset.py

@@ -1,160 +1,160 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-
-import ocnn
-from ocnn.octree import Octree, Points
-
-
-__all__ = ['Transform', 'CollateBatch']
-classes = __all__
-
-
-class Transform:
-  r''' A boilerplate class which transforms an input data for :obj:`ocnn`.
-  The input data is first converted to :class:`Points`, then randomly transformed
-  (if enabled), and converted to an :class:`Octree`.
-
-  Args:
-    depth (int): The octree depth.
-    full_depth (int): The octree layers with a depth small than
-        :attr:`full_depth` are forced to be full.
-    distort (bool): If true, performs the data augmentation.
-    angle (list): A list of 3 float values to generate random rotation angles.
-    interval (list): A list of 3 float values to represent the interval of
-        rotation angles.
-    scale (float): The maximum relative scale factor.
-    uniform (bool): If true, performs uniform scaling.
-    jittor (float): The maximum jitter values.
-    orient_normal (str): Orient point normals along the specified axis, which is
-        useful when normals are not oriented.
-  '''
-
-  def __init__(self, depth: int, full_depth: int, distort: bool, angle: list,
-               interval: list, scale: float, uniform: bool, jitter: float,
-               flip: list, orient_normal: str = '', **kwargs):
-    super().__init__()
-
-    # for octree building
-    self.depth = depth
-    self.full_depth = full_depth
-
-    # for data augmentation
-    self.distort = distort
-    self.angle = angle
-    self.interval = interval
-    self.scale = scale
-    self.uniform = uniform
-    self.jitter = jitter
-    self.flip = flip
-
-    # for other transformations
-    self.orient_normal = orient_normal
-
-  def __call__(self, sample: dict, idx: int):
-    r''''''
-
-    output = self.preprocess(sample, idx)
-    output = self.transform(output, idx)
-    output['octree'] = self.points2octree(output['points'])
-    return output
-
-  def preprocess(self, sample: dict, idx: int):
-    r''' Transforms :attr:`sample` to :class:`Points` and performs some specific
-    transformations, like normalization.
-    '''
-
-    xyz = torch.from_numpy(sample.pop('points'))
-    normals = torch.from_numpy(sample.pop('normals'))
-    sample['points'] = Points(xyz, normals)
-    return sample
-
-  def transform(self, sample: dict, idx: int):
-    r''' Applies the general transformations provided by :obj:`ocnn`.
-    '''
-
-    # The augmentations including rotation, scaling, and jittering.
-    points = sample['points']
-    if self.distort:
-      rng_angle, rng_scale, rng_jitter, rnd_flip = self.rnd_parameters()
-      points.flip(rnd_flip)
-      points.rotate(rng_angle)
-      points.translate(rng_jitter)
-      points.scale(rng_scale)
-
-    if self.orient_normal:
-      points.orient_normal(self.orient_normal)
-
-    # !!! NOTE: Clip the point cloud to [-1, 1] before building the octree
-    inbox_mask = points.clip(min=-1, max=1)
-    sample.update({'points': points, 'inbox_mask': inbox_mask})
-    return sample
-
-  def points2octree(self, points: Points):
-    r''' Converts the input :attr:`points` to an octree.
-    '''
-
-    octree = Octree(self.depth, self.full_depth)
-    octree.build_octree(points)
-    return octree
-
-  def rnd_parameters(self):
-    r''' Generates random parameters for data augmentation.
-    '''
-
-    rnd_angle = [None] * 3
-    for i in range(3):
-      rot_num = self.angle[i] // self.interval[i]
-      rnd = torch.randint(low=-rot_num, high=rot_num+1, size=(1,))
-      rnd_angle[i] = rnd * self.interval[i] * (3.14159265 / 180.0)
-    rnd_angle = torch.cat(rnd_angle)
-
-    rnd_scale = torch.rand(3) * (2 * self.scale) - self.scale + 1.0
-    if self.uniform:
-      rnd_scale[1] = rnd_scale[0]
-      rnd_scale[2] = rnd_scale[0]
-
-    rnd_flip = ''
-    for i, c in enumerate('xyz'):
-      if torch.rand([1]) < self.flip[i]:
-        rnd_flip = rnd_flip + c
-
-    rnd_jitter = torch.rand(3) * (2 * self.jitter) - self.jitter
-    return rnd_angle, rnd_scale, rnd_jitter, rnd_flip
-
-
-class CollateBatch:
-  r''' Merge a list of octrees and points into a batch.
-  '''
-
-  def __init__(self, merge_points: bool = False):
-    self.merge_points = merge_points
-
-  def __call__(self, batch: list):
-    assert type(batch) == list
-
-    outputs = {}
-    for key in batch[0].keys():
-      outputs[key] = [b[key] for b in batch]
-
-      # Merge a batch of octrees into one super octree
-      if 'octree' in key:
-        octree = ocnn.octree.merge_octrees(outputs[key])
-        # NOTE: remember to construct the neighbor indices
-        octree.construct_all_neigh()
-        outputs[key] = octree
-
-      # Merge a batch of points
-      if 'points' in key and self.merge_points:
-        outputs[key] = ocnn.octree.merge_points(outputs[key])
-
-      # Convert the labels to a Tensor
-      if 'label' in key:
-        outputs['label'] = torch.tensor(outputs[key])
-
-    return outputs
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+
+import ocnn
+from ocnn.octree import Octree, Points
+
+
+__all__ = ['Transform', 'CollateBatch']
+classes = __all__
+
+
+class Transform:
+  r''' A boilerplate class which transforms an input data for :obj:`ocnn`.
+  The input data is first converted to :class:`Points`, then randomly transformed
+  (if enabled), and converted to an :class:`Octree`.
+
+  Args:
+    depth (int): The octree depth.
+    full_depth (int): The octree layers with a depth small than
+        :attr:`full_depth` are forced to be full.
+    distort (bool): If true, performs the data augmentation.
+    angle (list): A list of 3 float values to generate random rotation angles.
+    interval (list): A list of 3 float values to represent the interval of
+        rotation angles.
+    scale (float): The maximum relative scale factor.
+    uniform (bool): If true, performs uniform scaling.
+    jittor (float): The maximum jitter values.
+    orient_normal (str): Orient point normals along the specified axis, which is
+        useful when normals are not oriented.
+  '''
+
+  def __init__(self, depth: int, full_depth: int, distort: bool, angle: list,
+               interval: list, scale: float, uniform: bool, jitter: float,
+               flip: list, orient_normal: str = '', **kwargs):
+    super().__init__()
+
+    # for octree building
+    self.depth = depth
+    self.full_depth = full_depth
+
+    # for data augmentation
+    self.distort = distort
+    self.angle = angle
+    self.interval = interval
+    self.scale = scale
+    self.uniform = uniform
+    self.jitter = jitter
+    self.flip = flip
+
+    # for other transformations
+    self.orient_normal = orient_normal
+
+  def __call__(self, sample: dict, idx: int):
+    r''''''
+
+    output = self.preprocess(sample, idx)
+    output = self.transform(output, idx)
+    output['octree'] = self.points2octree(output['points'])
+    return output
+
+  def preprocess(self, sample: dict, idx: int):
+    r''' Transforms :attr:`sample` to :class:`Points` and performs some specific
+    transformations, like normalization.
+    '''
+
+    xyz = torch.from_numpy(sample.pop('points'))
+    normals = torch.from_numpy(sample.pop('normals'))
+    sample['points'] = Points(xyz, normals)
+    return sample
+
+  def transform(self, sample: dict, idx: int):
+    r''' Applies the general transformations provided by :obj:`ocnn`.
+    '''
+
+    # The augmentations including rotation, scaling, and jittering.
+    points = sample['points']
+    if self.distort:
+      rng_angle, rng_scale, rng_jitter, rnd_flip = self.rnd_parameters()
+      points.flip(rnd_flip)
+      points.rotate(rng_angle)
+      points.translate(rng_jitter)
+      points.scale(rng_scale)
+
+    if self.orient_normal:
+      points.orient_normal(self.orient_normal)
+
+    # !!! NOTE: Clip the point cloud to [-1, 1] before building the octree
+    inbox_mask = points.clip(min=-1, max=1)
+    sample.update({'points': points, 'inbox_mask': inbox_mask})
+    return sample
+
+  def points2octree(self, points: Points):
+    r''' Converts the input :attr:`points` to an octree.
+    '''
+
+    octree = Octree(self.depth, self.full_depth)
+    octree.build_octree(points)
+    return octree
+
+  def rnd_parameters(self):
+    r''' Generates random parameters for data augmentation.
+    '''
+
+    rnd_angle = [None] * 3
+    for i in range(3):
+      rot_num = self.angle[i] // self.interval[i]
+      rnd = torch.randint(low=-rot_num, high=rot_num+1, size=(1,))
+      rnd_angle[i] = rnd * self.interval[i] * (3.14159265 / 180.0)
+    rnd_angle = torch.cat(rnd_angle)
+
+    rnd_scale = torch.rand(3) * (2 * self.scale) - self.scale + 1.0
+    if self.uniform:
+      rnd_scale[1] = rnd_scale[0]
+      rnd_scale[2] = rnd_scale[0]
+
+    rnd_flip = ''
+    for i, c in enumerate('xyz'):
+      if torch.rand([1]) < self.flip[i]:
+        rnd_flip = rnd_flip + c
+
+    rnd_jitter = torch.rand(3) * (2 * self.jitter) - self.jitter
+    return rnd_angle, rnd_scale, rnd_jitter, rnd_flip
+
+
+class CollateBatch:
+  r''' Merge a list of octrees and points into a batch.
+  '''
+
+  def __init__(self, merge_points: bool = False):
+    self.merge_points = merge_points
+
+  def __call__(self, batch: list):
+    assert type(batch) == list
+
+    outputs = {}
+    for key in batch[0].keys():
+      outputs[key] = [b[key] for b in batch]
+
+      # Merge a batch of octrees into one super octree
+      if 'octree' in key:
+        octree = ocnn.octree.merge_octrees(outputs[key])
+        # NOTE: remember to construct the neighbor indices
+        octree.construct_all_neigh()
+        outputs[key] = octree
+
+      # Merge a batch of points
+      if 'points' in key and self.merge_points:
+        outputs[key] = ocnn.octree.merge_points(outputs[key])
+
+      # Convert the labels to a Tensor
+      if 'label' in key:
+        outputs['label'] = torch.tensor(outputs[key])
+
+    return outputs

{ocnn-2.2.5 → ocnn-2.2.7}/ocnn/models/__init__.py

@@ -1,29 +1,29 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-from .lenet import LeNet
-from .resnet import ResNet
-from .segnet import SegNet
-from .unet import UNet
-from .hrnet import HRNet
-from .autoencoder import AutoEncoder
-from .ounet import OUNet
-from .image2shape import Image2Shape
-
-
-__all__ = [
-    'LeNet',
-    'ResNet',
-    'SegNet',
-    'UNet',
-    'HRNet',
-    'AutoEncoder',
-    'OUNet',
-    'Image2Shape',
-]
-
-classes = __all__
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+from .lenet import LeNet
+from .resnet import ResNet
+from .segnet import SegNet
+from .unet import UNet
+from .hrnet import HRNet
+from .autoencoder import AutoEncoder
+from .ounet import OUNet
+from .image2shape import Image2Shape
+
+
+__all__ = [
+    'LeNet',
+    'ResNet',
+    'SegNet',
+    'UNet',
+    'HRNet',
+    'AutoEncoder',
+    'OUNet',
+    'Image2Shape',
+]
+
+classes = __all__