ocnn 2.2.7__py3-none-any.whl → 2.3.0__py3-none-any.whl

ocnn/octree/octree.py

@@ -43,11 +43,13 @@ class Octree:
   def __init__(self, depth: int, full_depth: int = 2, batch_size: int = 1,
                device: Union[torch.device, str] = 'cpu', **kwargs):
     super().__init__()
+    # configurations for initialization
     self.depth = depth
     self.full_depth = full_depth
     self.batch_size = batch_size
     self.device = device
 
+    # properties after building the octree
     self.reset()
 
   def reset(self):
@@ -63,12 +65,18 @@ class Octree:
     self.normals = [None] * num
     self.points = [None] * num
 
+    # self.nempty_masks, self.nempty_indices and self.nempty_neighs are
+    # for handling of non-empty nodes and are constructed on demand
+    self.nempty_masks = [None] * num
+    self.nempty_indices = [None] * num
+    self.nempty_neighs = [None] * num
+
     # octree node numbers in each octree layers.
-    # TODO: decide whether to settle them to 'gpu' or not?
+    # These are small 1-D tensors; just keep them on CPUs
     self.nnum = torch.zeros(num, dtype=torch.long)
     self.nnum_nempty = torch.zeros(num, dtype=torch.long)
 
-    # the following properties are valid after `merge_octrees`.
+    # the following properties are only valid after `merge_octrees`.
     # TODO: make them valid after `octree_grow`, `octree_split` and `build_octree`
     batch_size = self.batch_size
     self.batch_nnum = torch.zeros(num, batch_size, dtype=torch.long)
@@ -107,8 +115,8 @@ class Octree:
 
     key = self.keys[depth]
     if nempty:
-      mask = self.nempty_mask(depth)
-      key = key[mask]
+      idx = self.nempty_index(depth)
+      key = key[idx]
     return key
 
   def xyzb(self, depth: int, nempty: bool = False):
@@ -132,19 +140,63 @@ class Octree:
 
     batch_id = self.keys[depth] >> 48
     if nempty:
-      mask = self.nempty_mask(depth)
-      batch_id = batch_id[mask]
+      idx = self.nempty_index(depth)
+      batch_id = batch_id[idx]
     return batch_id
 
-  def nempty_mask(self, depth: int):
+  def nempty_mask(self, depth: int, reset: bool = False):
     r''' Returns a binary mask which indicates whether the cooreponding octree
     node is empty or not.
 
     Args:
       depth (int): The depth of the octree.
+      reset (bool): If True, recomputes the mask.
     '''
 
-    return self.children[depth] >= 0
+    if self.nempty_masks[depth] is None or reset:
+      self.nempty_masks[depth] = self.children[depth] >= 0
+    return self.nempty_masks[depth]
+
+  def nempty_index(self, depth: int, reset: bool = False):
+    r''' Returns the indices of non-empty octree nodes.
+
+    Args:
+      depth (int): The depth of the octree.
+      reset (bool): If True, recomputes the indices.
+    '''
+
+    if self.nempty_indices[depth] is None or reset:
+      mask = self.nempty_mask(depth)
+      rng = torch.arange(mask.shape[0], device=mask.device, dtype=torch.long)
+      self.nempty_indices[depth] = rng[mask]
+    return self.nempty_indices[depth]
+
+  def nempty_neigh(self, depth: int, reset: bool = False):
+    r''' Returns the neighborhoods of non-empty octree nodes.
+    Args:
+
+      depth (int): The depth of the octree.
+      reset (bool): If True, recomputes the neighborhoods.
+    '''
+    if self.nempty_neighs[depth] is None or reset:
+      neigh = self.neighs[depth]
+      idx = self.nempty_index(depth)
+      neigh = self.remap_nempty_neigh(neigh[idx], depth)
+      self.nempty_neighs[depth] = neigh
+    return self.nempty_neighs[depth]
+
+  def remap_nempty_neigh(self, neigh: torch.Tensor, depth: int):
+    r''' Remaps the neighborhood indices to the non-empty octree nodes.
+
+    Args:
+      neigh (torch.Tensor): The input neighborhoods with shape :obj:`(N, 27)`.
+      depth (int): The depth of the octree.
+    '''
+
+    valid = neigh >= 0
+    child = self.children[depth]
+    neigh[valid] = child[neigh[valid]].long()  # remap the index
+    return neigh
 
   def build_octree(self, point_cloud: Points):
     r''' Builds an octree from a point cloud.
@@ -225,6 +277,10 @@ class Octree:
       features = scatter_add(point_cloud.features, idx, dim=0)
       self.features[d] = features / counts.unsqueeze(1)
 
+    # reset nempty_masks and nempty_indices, which will be updated on demand
+    self.nempty_masks = [None] * (self.depth + 1)
+    self.nempty_indices = [None] * (self.depth + 1)
+    self.nempty_neighs = [None] * (self.depth + 1)
     return idx
 
   def octree_grow_full(self, depth: int, update_neigh: bool = True):
@@ -240,18 +296,22 @@ class Octree:
 
     # node number
     num = 1 << (3 * depth)
-    self.nnum[depth] = num * self.batch_size
-    self.nnum_nempty[depth] = num * self.batch_size
+    batch_size = self.batch_size
+    self.nnum[depth] = num * batch_size
+    self.nnum_nempty[depth] = num * batch_size
 
     # update key
     key = torch.arange(num, dtype=torch.long, device=self.device)
-    bs = torch.arange(self.batch_size, dtype=torch.long, device=self.device)
+    bs = torch.arange(batch_size, dtype=torch.long, device=self.device)
     key = key.unsqueeze(0) | (bs.unsqueeze(1) << 48)
     self.keys[depth] = key.view(-1)
 
     # update children
     self.children[depth] = torch.arange(
-        num * self.batch_size, dtype=torch.int32, device=self.device)
+        num * batch_size, dtype=torch.int32, device=self.device)
+
+    # nempty_masks, nempty_indices, and nempty_neighs
+    # need not be reset for full octrees
 
     # update neigh if needed
     if update_neigh:
@@ -281,6 +341,12 @@ class Octree:
     self.children[depth] = children.int()
     self.nnum_nempty[depth] = nnum_nempty
 
+    # reset nempty_masks, nempty_indices, and nempty_neighs as they depend on
+    # children[depth] and are invalid now
+    self.nempty_masks[depth] = None
+    self.nempty_indices[depth] = None
+    self.nempty_neighs[depth] = None
+
   def octree_grow(self, depth: int, update_neigh: bool = True):
     r''' Grows the octree and updates the relevant properties. And in most
     cases, call :func:`Octree.octree_split` to update the splitting status of
@@ -301,6 +367,9 @@ class Octree:
       self.features.append(None)
       self.normals.append(None)
       self.points.append(None)
+      self.nempty_masks.append(None)
+      self.nempty_indices.append(None)
+      self.nempty_neighs.append(None)
       zero = torch.zeros(1, dtype=torch.long)
       self.nnum = torch.cat([self.nnum, zero])
       self.nnum_nempty = torch.cat([self.nnum_nempty, zero])
@@ -342,7 +411,7 @@ class Octree:
       key = torch.arange(nnum, dtype=torch.long, device=device)
       x, y, z, _ = key2xyz(key, depth)
       xyz = torch.stack([x, y, z], dim=-1)  # (N,  3)
-      grid = range_grid(-1, 1, device)   # (27, 3)
+      grid = range_grid(-1, 1, device)      # (27, 3)
       xyz = xyz.unsqueeze(1) + grid         # (N, 27, 3)
       xyz = xyz.view(-1, 3)                 # (N*27, 3)
       neigh = xyz2key(xyz[:, 0], xyz[:, 1], xyz[:, 2], depth=depth)
@@ -400,16 +469,18 @@ class Octree:
     '''
 
     key = self.key(depth, nempty)
-    # `torch.bucketize` is similar to `torch.searchsorted`.
-    # I choose `torch.bucketize` here because it has fewer dimension checks,
-    # resulting in slightly better performance according to the docs of
-    # pytorch-1.9.1, since `key` is always 1-D sorted sequence.
+    idx = torch.searchsorted(key, query)
+
+    # `torch.bucketize` can also be used here; it is similar to
+    # `torch.searchsorted`, and it has fewer dimension checks, resulting in
+    # slightly better performance for 1D sorted sequences according to the docs
+    # of pytorch-1.9.1. `key` is always a 1D sorted sequence.
     # https://pytorch.org/docs/1.9.1/generated/torch.searchsorted.html
-    idx = torch.bucketize(query, key)
+    # idx = torch.bucketize(query, key)
 
     valid = idx < key.shape[0]      # valid if in-bound
-    found = key[idx[valid]] == query[valid]
-    valid[valid.clone()] = found    # valid if found
+    vi = torch.arange(query.shape[0], device=query.device)[valid]
+    valid[vi] = key[idx[vi]] == query[vi]  # valid if found
     idx[valid.logical_not()] = -1   # set to -1 if invalid
     return idx
 
@@ -428,22 +499,18 @@ class Octree:
           octree nodes.
     '''
 
-    if stride == 1:
+    if stride == 1 and not nempty:
       neigh = self.neighs[depth]
-    elif stride == 2:
-      # clone neigh to avoid self.neigh[depth] being modified
+    elif stride == 2 and not nempty:
       neigh = self.neighs[depth][::8].clone()
+    elif stride == 1 and nempty:
+      neigh = self.nempty_neigh(depth)
+    elif stride == 2 and nempty:
+      neigh = self.neighs[depth][::8].clone()
+      neigh = self.remap_nempty_neigh(neigh, depth)
     else:
       raise ValueError('Unsupported stride {}'.format(stride))
 
-    if nempty:
-      child = self.children[depth]
-      if stride == 1:
-        nempty_node = child >= 0
-        neigh = neigh[nempty_node]
-      valid = neigh >= 0
-      neigh[valid] = child[neigh[valid]].long()  # remap the index
-
     if kernel == '333':
       return neigh
     elif kernel in self.lut_kernel:
@@ -548,15 +615,23 @@ class Octree:
       return [p.to(device, non_blocking=non_blocking)
               if isinstance(p, torch.Tensor) else None for p in prop]
 
-    # Construct a new Octree on the specified device
-    octree = Octree(self.depth, self.full_depth, self.batch_size, device)
+    # Construct a new Octree on the specified device.
+    # During the initialization, self.device is used to set up the new Octree;
+    # the look-up tables (including self.lut_kernel, self.lut_parent, and
+    # self.lut_child), will be already created on the correct device.
+    octree = Octree.init_like(self, device)
+
+    # Move all the other properties to the specified device
     octree.keys = list_to_device(self.keys)
     octree.children = list_to_device(self.children)
     octree.neighs = list_to_device(self.neighs)
     octree.features = list_to_device(self.features)
     octree.normals = list_to_device(self.normals)
     octree.points = list_to_device(self.points)
-    octree.nnum = self.nnum.clone()  # TODO: whether to move nnum to the self.device?
+
+    # The following are small tensors, keep them on CPU to avoid frequent device
+    # switching, so just clone them.
+    octree.nnum = self.nnum.clone()
     octree.nnum_nempty = self.nnum_nempty.clone()
     octree.batch_nnum = self.batch_nnum.clone()
     octree.batch_nnum_nempty = self.batch_nnum_nempty.clone()
@@ -572,90 +647,124 @@ class Octree:
 
     return self.to('cpu')
 
+  def merge_octrees(self, octrees: List['Octree']):
+    r''' Merges a list of octrees into one batch.
 
-def merge_octrees(octrees: List['Octree']):
-  r''' Merges a list of octrees into one batch.
+    Args:
+      octrees (List[Octree]): A list of octrees to merge.
 
-  Args:
-    octrees (List[Octree]): A list of octrees to merge.
+    Returns:
+      Octree: The merged octree.
+    '''
+
+    # init and check
+    batch_size = len(octrees)
+    self.batch_size = batch_size
+    for i in range(1, batch_size):
+      condition = (octrees[i].depth == self.depth and
+                   octrees[i].full_depth == self.full_depth and
+                   octrees[i].device == self.device)
+      assert condition, 'The check of merge_octrees failed'
+
+    # node num
+    batch_nnum = torch.stack(
+        [octrees[i].nnum for i in range(batch_size)], dim=1)
+    batch_nnum_nempty = torch.stack(
+        [octrees[i].nnum_nempty for i in range(batch_size)], dim=1)
+    self.nnum = torch.sum(batch_nnum, dim=1)
+    self.nnum_nempty = torch.sum(batch_nnum_nempty, dim=1)
+    self.batch_nnum = batch_nnum
+    self.batch_nnum_nempty = batch_nnum_nempty
+    nnum_cum = cumsum(batch_nnum_nempty, dim=1, exclusive=True)
+
+    # merge octre properties
+    for d in range(self.depth + 1):
+      # key
+      keys = [None] * batch_size
+      for i in range(batch_size):
+        key = octrees[i].keys[d] & ((1 << 48) - 1)  # clear the highest bits
+        keys[i] = key | (i << 48)
+      self.keys[d] = torch.cat(keys, dim=0)
+
+      # children
+      children = [None] * batch_size
+      for i in range(batch_size):
+        # !! `clone` is used here to avoid modifying the original octrees
+        child = octrees[i].children[d].clone()
+        mask = child >= 0
+        child[mask] = child[mask] + nnum_cum[d, i]
+        children[i] = child
+      self.children[d] = torch.cat(children, dim=0)
+
+      # features
+      if octrees[0].features[d] is not None and d == self.depth:
+        features = [octrees[i].features[d] for i in range(batch_size)]
+        self.features[d] = torch.cat(features, dim=0)
+
+      # normals
+      if octrees[0].normals[d] is not None and d == self.depth:
+        normals = [octrees[i].normals[d] for i in range(batch_size)]
+        self.normals[d] = torch.cat(normals, dim=0)
+
+      # points
+      if octrees[0].points[d] is not None and d == self.depth:
+        points = [octrees[i].points[d] for i in range(batch_size)]
+        self.points[d] = torch.cat(points, dim=0)
+
+    return self
+
+  @classmethod
+  def init_like(cls, octree: 'Octree', device: Union[torch.device, str, None] = None):
+    r''' Initializes the octree like another octree.
+
+    Args:
+      octree (Octree): The reference octree.
+      device (torch.device or str): The device to use for computation.
+    '''
+
+    device = device if device is not None else octree.device
+    return cls(depth=octree.depth, full_depth=octree.full_depth,
+               batch_size=octree.batch_size, device=device)
+
+  @classmethod
+  def init_octree(cls, depth: int, full_depth: int = 2, batch_size: int = 1,
+                  device: Union[torch.device, str] = 'cpu'):
+    r'''
+    Initializes an octree to :attr:`full_depth`.
 
-  Returns:
-    Octree: The merged octree.
+    Args:
+      depth (int): The depth of the octree.
+      full_depth (int): The octree layers with a depth small than
+          :attr:`full_depth` are forced to be full.
+      batch_size (int, optional): The batch size.
+      device (torch.device or str): The device to use for computation.
+
+    Returns:
+      Octree: The initialized Octree object.
+    '''
+
+    octree = cls(depth, full_depth, batch_size, device)
+    for d in range(full_depth + 1):
+      octree.octree_grow_full(depth=d)
+    return octree
+
+
+def merge_octrees(octrees: List['Octree']):
+  r''' A wrapper of :meth:`Octree.merge_octrees`.
+
+  .. deprecated:: 2.2.7
+    Use :meth:`Octree.merge_octrees` instead.
   '''
 
-  # init and check
-  octree = Octree(depth=octrees[0].depth, full_depth=octrees[0].full_depth,
-                  batch_size=len(octrees), device=octrees[0].device)
-  for i in range(1, octree.batch_size):
-    condition = (octrees[i].depth == octree.depth and
-                 octrees[i].full_depth == octree.full_depth and
-                 octrees[i].device == octree.device)
-    assert condition, 'The check of merge_octrees failed'
-
-  # node num
-  batch_nnum = torch.stack(
-      [octrees[i].nnum for i in range(octree.batch_size)], dim=1)
-  batch_nnum_nempty = torch.stack(
-      [octrees[i].nnum_nempty for i in range(octree.batch_size)], dim=1)
-  octree.nnum = torch.sum(batch_nnum, dim=1)
-  octree.nnum_nempty = torch.sum(batch_nnum_nempty, dim=1)
-  octree.batch_nnum = batch_nnum
-  octree.batch_nnum_nempty = batch_nnum_nempty
-  nnum_cum = cumsum(batch_nnum_nempty, dim=1, exclusive=True)
-
-  # merge octre properties
-  for d in range(octree.depth+1):
-    # key
-    keys = [None] * octree.batch_size
-    for i in range(octree.batch_size):
-      key = octrees[i].keys[d] & ((1 << 48) - 1)  # clear the highest bits
-      keys[i] = key | (i << 48)
-    octree.keys[d] = torch.cat(keys, dim=0)
-
-    # children
-    children = [None] * octree.batch_size
-    for i in range(octree.batch_size):
-      child = octrees[i].children[d].clone()  # !! `clone` is used here to avoid
-      mask = child >= 0                       # !! modifying the original octrees
-      child[mask] = child[mask] + nnum_cum[d, i]
-      children[i] = child
-    octree.children[d] = torch.cat(children, dim=0)
-
-    # features
-    if octrees[0].features[d] is not None and d == octree.depth:
-      features = [octrees[i].features[d] for i in range(octree.batch_size)]
-      octree.features[d] = torch.cat(features, dim=0)
-
-    # normals
-    if octrees[0].normals[d] is not None and d == octree.depth:
-      normals = [octrees[i].normals[d] for i in range(octree.batch_size)]
-      octree.normals[d] = torch.cat(normals, dim=0)
-
-    # points
-    if octrees[0].points[d] is not None and d == octree.depth:
-      points = [octrees[i].points[d] for i in range(octree.batch_size)]
-      octree.points[d] = torch.cat(points, dim=0)
-
-  return octree
+  return Octree.init_like(octrees[0]).merge_octrees(octrees)
 
 
 def init_octree(depth: int, full_depth: int = 2, batch_size: int = 1,
                 device: Union[torch.device, str] = 'cpu'):
-  r'''
-  Initializes an octree to :attr:`full_depth`.
-
-  Args:
-    depth (int): The depth of the octree.
-    full_depth (int): The octree layers with a depth small than
-        :attr:`full_depth` are forced to be full.
-    batch_size (int, optional): The batch size.
-    device (torch.device or str): The device to use for computation.
+  r''' A wrapper of :meth:`Octree.init_octree`.
 
-  Returns:
-    Octree: The initialized Octree object.
+  .. deprecated:: 2.2.7
+    Use :meth:`Octree.init_octree` instead.
   '''
 
-  octree = Octree(depth, full_depth, batch_size, device)
-  for d in range(full_depth+1):
-    octree.octree_grow_full(depth=d)
-  return octree
+  return Octree.init_octree(depth, full_depth, batch_size, device)

ocnn/octree/points.py

@@ -8,6 +8,7 @@
 import torch
 import numpy as np
 from typing import Optional, Union, List
+from ocnn.utils import cumsum
 
 
 class Points:
@@ -66,6 +67,7 @@ class Points:
       if self.batch_id.dim() == 1:
         self.batch_id = self.batch_id.unsqueeze(1)
       assert self.batch_id.size(1) == 1
+      assert self.batch_size == self.batch_id.max().item() + 1
 
   @property
   def npt(self):
@@ -165,26 +167,23 @@ class Points:
     '''
 
     mask = self.inbox_mask(min + esp, max - esp)
-    tmp = self.__getitem__(mask)
-    self.__dict__.update(tmp.__dict__)
+    self.copy_from(self[mask])
     return mask
 
-  def __getitem__(self, mask: torch.Tensor):
-    r''' Slices the point cloud according a given :attr:`mask`.
+  def __getitem__(self, idx):
+    r''' Slices the point cloud according a given :attr:`idx`.
     '''
 
-    dummy_pts = torch.zeros(1, 3, device=self.device)
-    out = Points(dummy_pts, batch_size=self.batch_size)
-
-    out.points = self.points[mask]
+    out = self.init_points(self.device, self.batch_size)
+    out.points = self.points[idx]
     if self.normals is not None:
-      out.normals = self.normals[mask]
+      out.normals = self.normals[idx]
     if self.features is not None:
-      out.features = self.features[mask]
+      out.features = self.features[idx]
     if self.labels is not None:
-      out.labels = self.labels[mask]
+      out.labels = self.labels[idx]
     if self.batch_id is not None:
-      out.batch_id = self.batch_id[mask]
+      out.batch_id = self.batch_id[idx]
     return out
 
   def inbox_mask(self, bbmin: Union[float, torch.Tensor] = -1.0,
@@ -239,8 +238,7 @@ class Points:
       return self
 
     # Construct a new Points on the specified device
-    points = Points(torch.zeros(1, 3, device=device))
-    points.batch_size = self.batch_size
+    points = self.init_points(device, self.batch_size)
     points.batch_npt = self.batch_npt
     points.points = self.points.to(device, non_blocking=non_blocking)
     if self.normals is not None:
@@ -295,29 +293,92 @@ class Points:
     else:
       raise ValueError
 
+  def copy_from(self, points: 'Points'):
+    r''' Shallow copy from another Points.
+    '''
+
+    self.points = points.points
+    self.normals = points.normals
+    self.features = points.features
+    self.labels = points.labels
+    self.batch_id = points.batch_id
+    self.batch_size = points.batch_size
+    self.device = points.device
+    self.batch_npt = points.batch_npt
+
+  def merge_points(self, points: List['Points'], update_batch_info: bool = True):
+    r''' Merges a list of points into one batch.
+
+    Args:
+      points (List[Octree]): A list of points to merge. The batch size of each
+          points in the list is assumed to be 1, and the :obj:`batch_size`,
+          :obj:`batch_id`, and :obj:`batch_npt` in the points are ignored.
+    '''
+
+    self.points = torch.cat([p.points for p in points], dim=0)
+    if points[0].normals is not None:
+      self.normals = torch.cat([p.normals for p in points], dim=0)
+    if points[0].features is not None:
+      self.features = torch.cat([p.features for p in points], dim=0)
+    if points[0].labels is not None:
+      self.labels = torch.cat([p.labels for p in points], dim=0)
+    self.device = points[0].device
+
+    if update_batch_info:
+      self.batch_size = len(points)
+      self.batch_npt = torch.Tensor([p.npt for p in points]).long()
+      self.batch_id = torch.cat([p.points.new_full((p.npt, 1), i)
+                                 for i, p in enumerate(points)], dim=0)
+    return self
+
+  def split_points(self):
+    r''' Splits the batched points into a list of Points.
+    '''
+
+    if self.batch_npt is None:
+      self.batch_npt = torch.bincount(
+          self.batch_id.squeeze(), minlength=self.batch_size)
+
+    outs = []
+    cs = cumsum(self.batch_npt, dim=0, exclusive=True)
+    for i in range(self.batch_size):
+      rng = range(cs[i], cs[i+1])
+      out = Points.init_points(self.device, batch_size=1)
+      out.points = self.points[rng]
+      if self.normals is not None:
+        out.normals = self.normals[rng]
+      if self.features is not None:
+        out.features = self.features[rng]
+      if self.labels is not None:
+        out.labels = self.labels[rng]
+      outs.append(out)
+    return outs
+
+  @classmethod
+  def init_points(cls, device: Union[torch.device, str, None] = None,
+                  batch_size: int = 1):
+    r''' Initialzes a Points object with dummy data on a specified device.
+
+    Args:
+      device (torch.device or str or None): The device of the Points. If
+          :obj:`None`, the device is set to :obj:`cpu`.
+      batch_size (int): The batch size.
+    '''
+
+    points = torch.zeros(batch_size, 3, device=device)
+    batch_id = (torch.arange(batch_size, device=device).unsqueeze(1)
+                if batch_size > 1 else None)
+    return cls(points, batch_size=batch_size, batch_id=batch_id)
+
 
 def merge_points(points: List['Points'], update_batch_info: bool = True):
-  r''' Merges a list of points into one batch.
+  r''' A wrapper of :meth:`Points.merge_points`.
 
-  Args:
-    points (List[Octree]): A list of points to merge. The batch size of each
-        points in the list is assumed to be 1, and the :obj:`batch_size`,
-        :obj:`batch_id`, and :obj:`batch_npt` in the points are ignored.
+  .. deprecated:: 2.2.7
+     Use :meth:`Points.merge_points` instead.
   '''
 
-  out = Points(torch.zeros(1, 3))
-  out.points = torch.cat([p.points for p in points], dim=0)
-  if points[0].normals is not None:
-    out.normals = torch.cat([p.normals for p in points], dim=0)
-  if points[0].features is not None:
-    out.features = torch.cat([p.features for p in points], dim=0)
-  if points[0].labels is not None:
-    out.labels = torch.cat([p.labels for p in points], dim=0)
-  out.device = points[0].device
-
-  if update_batch_info:
-    out.batch_size = len(points)
-    out.batch_npt = torch.Tensor([p.npt for p in points]).long()
-    out.batch_id = torch.cat([p.points.new_full((p.npt, 1), i)
-                              for i, p in enumerate(points)], dim=0)
+  assert len(points) > 0, 'The input points list is empty.'
+  out = Points.init_points(points[0].device, batch_size=len(points))
+  out.merge_points(points, update_batch_info)
   return out

{ocnn-2.2.7.dist-info → ocnn-2.3.0.dist-info}/METADATA

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: ocnn
-Version: 2.2.7
+Version: 2.3.0
 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
@@ -84,6 +83,13 @@ 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.
 
+## Updates
+
+- **2026.02.02**: Release `v2.3.0`, incorporating Triton to accelerate
+  octree-based sparse convolution in the upcoming release. OctreeConv is even
+  **2.5 times faster than the latest spconv**!
+- **2025.12.18**: Release `v2.2.8`, improving neighbor search efficiency.
+
 
 ## Key benefits of ocnn-pytorch
 
@@ -93,10 +99,9 @@ octrees to leverage the library's capabilities.
   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.
+  convolution frameworks. It is **even 2.5 times faster than the latest spconv
+  implementation**! Check the benchmark [code](test/benchmark_conv.py) and
+  [results](test/benchmark/results.png) for details. ✨
 
 ## Citation