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

ocnn/models/lenet.py

@@ -1,46 +1,46 @@
-# --------------------------------------------------------
-# 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
-
-
-class LeNet(torch.nn.Module):
-  r''' Octree-based LeNet for classification.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, stages: int,
-               nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.stages = stages
-    self.nempty = nempty
-    channels = [in_channels] + [2 ** max(i+7-stages, 2) for i in range(stages)]
-
-    self.convs = torch.nn.ModuleList([ocnn.modules.OctreeConvBnRelu(
-        channels[i], channels[i+1], nempty=nempty) for i in range(stages)])
-    self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
-        nempty) for _ in range(stages)])
-    self.octree2voxel = ocnn.nn.Octree2Voxel(self.nempty)
-    self.header = torch.nn.Sequential(
-        torch.nn.Dropout(p=0.5),                     # drop1
-        ocnn.modules.FcBnRelu(64 * 64, 128),         # fc1
-        torch.nn.Dropout(p=0.5),                     # drop2
-        torch.nn.Linear(128, out_channels))          # fc2
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    for i in range(self.stages):
-      d = depth - i
-      data = self.convs[i](data, octree, d)
-      data = self.pools[i](data, octree, d)
-    data = self.octree2voxel(data, octree, depth-self.stages)
-    data = self.header(data)
-    return data
+# --------------------------------------------------------
+# 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
+
+
+class LeNet(torch.nn.Module):
+  r''' Octree-based LeNet for classification.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, stages: int,
+               nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.stages = stages
+    self.nempty = nempty
+    channels = [in_channels] + [2 ** max(i+7-stages, 2) for i in range(stages)]
+
+    self.convs = torch.nn.ModuleList([ocnn.modules.OctreeConvBnRelu(
+        channels[i], channels[i+1], nempty=nempty) for i in range(stages)])
+    self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
+        nempty) for _ in range(stages)])
+    self.octree2voxel = ocnn.nn.Octree2Voxel(self.nempty)
+    self.header = torch.nn.Sequential(
+        torch.nn.Dropout(p=0.5),                     # drop1
+        ocnn.modules.FcBnRelu(64 * 64, 128),         # fc1
+        torch.nn.Dropout(p=0.5),                     # drop2
+        torch.nn.Linear(128, out_channels))          # fc2
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    for i in range(self.stages):
+      d = depth - i
+      data = self.convs[i](data, octree, d)
+      data = self.pools[i](data, octree, d)
+    data = self.octree2voxel(data, octree, depth-self.stages)
+    data = self.header(data)
+    return data

ocnn/models/ounet.py

@@ -1,94 +1,94 @@
-# --------------------------------------------------------
-# 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 ocnn
-import torch
-import torch.nn
-
-from ocnn.octree import Octree
-from ocnn.models.autoencoder import AutoEncoder
-
-
-class OUNet(AutoEncoder):
-
-  def __init__(self, channel_in: int, channel_out: int, depth: int,
-               full_depth: int = 2, feature: str = 'ND'):
-    super().__init__(channel_in, channel_out, depth, full_depth, feature)
-    self.proj = None  # remove this module used in AutoEncoder
-
-  def encoder(self, octree):
-    r''' The encoder network for extracting heirarchy features.
-    '''
-
-    convs = dict()
-    depth, full_depth = self.depth, self.full_depth
-    data = octree.get_input_feature(self.feature, nempty=False)
-    assert data.size(1) == self.channel_in
-    convs[depth] = self.conv1(data, octree, depth)
-    for i, d in enumerate(range(depth, full_depth-1, -1)):
-      convs[d] = self.encoder_blks[i](convs[d], octree, d)
-      if d > full_depth:
-        convs[d-1] = self.downsample[i](convs[d], octree, d)
-    return convs
-
-  def decoder(self, convs: dict, octree_in: Octree, octree_out: Octree,
-              update_octree: bool = False):
-    r''' The decoder network for decode the octree.
-    '''
-
-    logits = dict()
-    deconv = convs[self.full_depth]
-    depth, full_depth = self.depth, self.full_depth
-    for i, d in enumerate(range(full_depth, depth + 1)):
-      if d > full_depth:
-        deconv = self.upsample[i-1](deconv, octree_out, d-1)
-        skip = ocnn.nn.octree_align(convs[d], octree_in, octree_out, d)
-        deconv = deconv + skip  # output-guided skip connections
-      deconv = self.decoder_blks[i](deconv, octree_out, d)
-
-      # predict the splitting label
-      logit = self.predict[i](deconv)
-      logits[d] = logit
-
-      # update the octree according to predicted labels
-      if update_octree:
-        split = logit.argmax(1).int()
-        octree_out.octree_split(split, d)
-        if d < depth:
-          octree_out.octree_grow(d + 1)
-
-      # predict the signal
-      if d == depth:
-        signal = self.header(deconv)
-        signal = torch.tanh(signal)
-        signal = ocnn.nn.octree_depad(signal, octree_out, depth)
-        if update_octree:
-          octree_out.features[depth] = signal
-
-    return {'logits': logits, 'signal': signal, 'octree_out': octree_out}
-
-  def init_octree(self, octree_in: Octree):
-    r''' Initialize a full octree for decoding.
-    '''
-
-    device = octree_in.device
-    batch_size = octree_in.batch_size
-    octree = Octree(self.depth, self.full_depth, batch_size, device)
-    for d in range(self.full_depth+1):
-      octree.octree_grow_full(depth=d)
-    return octree
-
-  def forward(self, octree_in, octree_out=None, update_octree: bool = False):
-    r''''''
-
-    if octree_out is None:
-      update_octree = True
-      octree_out = self.init_octree(octree_in)
-    convs = self.encoder(octree_in)
-    out = self.decoder(convs, octree_in, octree_out, update_octree)
-    return out
+# --------------------------------------------------------
+# 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 ocnn
+import torch
+import torch.nn
+
+from ocnn.octree import Octree
+from ocnn.models.autoencoder import AutoEncoder
+
+
+class OUNet(AutoEncoder):
+
+  def __init__(self, channel_in: int, channel_out: int, depth: int,
+               full_depth: int = 2, feature: str = 'ND'):
+    super().__init__(channel_in, channel_out, depth, full_depth, feature)
+    self.proj = None  # remove this module used in AutoEncoder
+
+  def encoder(self, octree):
+    r''' The encoder network for extracting heirarchy features.
+    '''
+
+    convs = dict()
+    depth, full_depth = self.depth, self.full_depth
+    data = octree.get_input_feature(self.feature, nempty=False)
+    assert data.size(1) == self.channel_in
+    convs[depth] = self.conv1(data, octree, depth)
+    for i, d in enumerate(range(depth, full_depth-1, -1)):
+      convs[d] = self.encoder_blks[i](convs[d], octree, d)
+      if d > full_depth:
+        convs[d-1] = self.downsample[i](convs[d], octree, d)
+    return convs
+
+  def decoder(self, convs: dict, octree_in: Octree, octree_out: Octree,
+              update_octree: bool = False):
+    r''' The decoder network for decode the octree.
+    '''
+
+    logits = dict()
+    deconv = convs[self.full_depth]
+    depth, full_depth = self.depth, self.full_depth
+    for i, d in enumerate(range(full_depth, depth + 1)):
+      if d > full_depth:
+        deconv = self.upsample[i-1](deconv, octree_out, d-1)
+        skip = ocnn.nn.octree_align(convs[d], octree_in, octree_out, d)
+        deconv = deconv + skip  # output-guided skip connections
+      deconv = self.decoder_blks[i](deconv, octree_out, d)
+
+      # predict the splitting label
+      logit = self.predict[i](deconv)
+      logits[d] = logit
+
+      # update the octree according to predicted labels
+      if update_octree:
+        split = logit.argmax(1).int()
+        octree_out.octree_split(split, d)
+        if d < depth:
+          octree_out.octree_grow(d + 1)
+
+      # predict the signal
+      if d == depth:
+        signal = self.header(deconv)
+        signal = torch.tanh(signal)
+        signal = ocnn.nn.octree_depad(signal, octree_out, depth)
+        if update_octree:
+          octree_out.features[depth] = signal
+
+    return {'logits': logits, 'signal': signal, 'octree_out': octree_out}
+
+  def init_octree(self, octree_in: Octree):
+    r''' Initialize a full octree for decoding.
+    '''
+
+    device = octree_in.device
+    batch_size = octree_in.batch_size
+    octree = Octree(self.depth, self.full_depth, batch_size, device)
+    for d in range(self.full_depth+1):
+      octree.octree_grow_full(depth=d)
+    return octree
+
+  def forward(self, octree_in, octree_out=None, update_octree: bool = False):
+    r''''''
+
+    if octree_out is None:
+      update_octree = True
+      octree_out = self.init_octree(octree_in)
+    convs = self.encoder(octree_in)
+    out = self.decoder(convs, octree_in, octree_out, update_octree)
+    return out

ocnn/models/resnet.py

@@ -1,53 +1,53 @@
-# --------------------------------------------------------
-# 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
-
-
-class ResNet(torch.nn.Module):
-  r''' Octree-based ResNet for classification.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, resblock_num: int,
-               stages: int, nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.resblk_num = resblock_num
-    self.stages = stages
-    self.nempty = nempty
-    channels = [2 ** max(i+9-stages, 2) for i in range(stages)]
-
-    self.conv1 = ocnn.modules.OctreeConvBnRelu(
-        in_channels, channels[0], nempty=nempty)
-    self.pool1 = ocnn.nn.OctreeMaxPool(nempty)
-    self.resblocks = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
-        channels[i], channels[i+1], resblock_num, nempty=nempty)
-        for i in range(stages-1)])
-    self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
-        nempty) for _ in range(stages-1)])
-    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
-    # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
-    self.header = torch.nn.Sequential(
-        ocnn.modules.FcBnRelu(channels[-1], 512),
-        torch.nn.Dropout(p=0.5),
-        torch.nn.Linear(512, out_channels))
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    data = self.conv1(data, octree, depth)
-    data = self.pool1(data, octree, depth)
-    for i in range(self.stages-1):
-      d = depth - i - 1
-      data = self.resblocks[i](data, octree, d)
-      data = self.pools[i](data, octree, d)
-    data = self.global_pool(data, octree, depth-self.stages)
-    data = self.header(data)
-    return data
+# --------------------------------------------------------
+# 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
+
+
+class ResNet(torch.nn.Module):
+  r''' Octree-based ResNet for classification.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, resblock_num: int,
+               stages: int, nempty: bool = False, dropout: float = 0.5):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.resblk_num = resblock_num
+    self.stages = stages
+    self.nempty = nempty
+    channels = [2 ** max(i+9-stages, 2) for i in range(stages)]
+
+    self.conv1 = ocnn.modules.OctreeConvBnRelu(
+        in_channels, channels[0], nempty=nempty)
+    self.pool1 = ocnn.nn.OctreeMaxPool(nempty)
+    self.resblocks = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
+        channels[i], channels[i+1], resblock_num, nempty=nempty)
+        for i in range(stages-1)])
+    self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
+        nempty) for _ in range(stages-1)])
+    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
+    # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
+    self.header = torch.nn.Sequential(
+        ocnn.modules.FcBnRelu(channels[-1], 512),
+        torch.nn.Dropout(p=dropout),
+        torch.nn.Linear(512, out_channels))
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    data = self.conv1(data, octree, depth)
+    data = self.pool1(data, octree, depth)
+    for i in range(self.stages-1):
+      d = depth - i - 1
+      data = self.resblocks[i](data, octree, d)
+      data = self.pools[i](data, octree, d)
+    data = self.global_pool(data, octree, depth-self.stages)
+    data = self.header(data)
+    return data

ocnn/models/segnet.py

@@ -1,72 +1,72 @@
-# --------------------------------------------------------
-# 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
-
-
-class SegNet(torch.nn.Module):
-  r''' Octree-based SegNet for segmentation.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, stages: int,
-               interp: str = 'linear', nempty: bool = False, **kwargs):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.stages = stages
-    self.nempty = nempty
-    return_indices = True
-
-    channels_stages = [2 ** max(i+8-stages, 2) for i in range(stages)]
-    channels = [in_channels] + channels_stages
-    self.convs = torch.nn.ModuleList(
-        [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
-         for i in range(stages)])
-    self.pools = torch.nn.ModuleList(
-        [ocnn.nn.OctreeMaxPool(nempty, return_indices) for _ in range(stages)])
-
-    self.bottleneck = ocnn.modules.OctreeConvBnRelu(channels[-1], channels[-1])
-
-    channels = channels_stages[::-1] + [channels_stages[0]]
-    self.deconvs = torch.nn.ModuleList(
-        [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
-         for i in range(0, stages)])
-    self.unpools = torch.nn.ModuleList(
-        [ocnn.nn.OctreeMaxUnpool(nempty) for _ in range(stages)])
-
-    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
-    self.header = torch.nn.Sequential(
-        ocnn.modules.Conv1x1BnRelu(channels[-1], 64),
-        ocnn.modules.Conv1x1(64, out_channels, use_bias=True))
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
-              query_pts: torch.Tensor):
-    r''''''
-
-    # encoder
-    indices = dict()
-    for i in range(self.stages):
-      d = depth - i
-      data = self.convs[i](data, octree, d)
-      data, indices[d] = self.pools[i](data, octree, d)
-
-    # bottleneck
-    data = self.bottleneck(data, octree, depth-self.stages)
-
-    # decoder
-    for i in range(self.stages):
-      d = depth - self.stages + i
-      data = self.unpools[i](data, indices[d + 1], octree, d)
-      data = self.deconvs[i](data, octree, d + 1)
-
-    # header
-    feature = self.octree_interp(data, octree, depth, query_pts)
-    logits = self.header(feature)
-
-    return logits
+# --------------------------------------------------------
+# 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
+
+
+class SegNet(torch.nn.Module):
+  r''' Octree-based SegNet for segmentation.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, stages: int,
+               interp: str = 'linear', nempty: bool = False, **kwargs):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.stages = stages
+    self.nempty = nempty
+    return_indices = True
+
+    channels_stages = [2 ** max(i+8-stages, 2) for i in range(stages)]
+    channels = [in_channels] + channels_stages
+    self.convs = torch.nn.ModuleList(
+        [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
+         for i in range(stages)])
+    self.pools = torch.nn.ModuleList(
+        [ocnn.nn.OctreeMaxPool(nempty, return_indices) for _ in range(stages)])
+
+    self.bottleneck = ocnn.modules.OctreeConvBnRelu(channels[-1], channels[-1])
+
+    channels = channels_stages[::-1] + [channels_stages[0]]
+    self.deconvs = torch.nn.ModuleList(
+        [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
+         for i in range(0, stages)])
+    self.unpools = torch.nn.ModuleList(
+        [ocnn.nn.OctreeMaxUnpool(nempty) for _ in range(stages)])
+
+    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
+    self.header = torch.nn.Sequential(
+        ocnn.modules.Conv1x1BnRelu(channels[-1], 64),
+        ocnn.modules.Conv1x1(64, out_channels, use_bias=True))
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
+              query_pts: torch.Tensor):
+    r''''''
+
+    # encoder
+    indices = dict()
+    for i in range(self.stages):
+      d = depth - i
+      data = self.convs[i](data, octree, d)
+      data, indices[d] = self.pools[i](data, octree, d)
+
+    # bottleneck
+    data = self.bottleneck(data, octree, depth-self.stages)
+
+    # decoder
+    for i in range(self.stages):
+      d = depth - self.stages + i
+      data = self.unpools[i](data, indices[d + 1], octree, d)
+      data = self.deconvs[i](data, octree, d + 1)
+
+    # header
+    feature = self.octree_interp(data, octree, depth, query_pts)
+    logits = self.header(feature)
+
+    return logits