ocnn 2.2.1__py3-none-any.whl → 2.2.2__py3-none-any.whl

ocnn/models/hrnet.py

@@ -1,192 +1,192 @@
-# --------------------------------------------------------
-# 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
-from typing import List
-
-import ocnn
-from ocnn.octree import Octree
-
-
-class Branches(torch.nn.Module):
-
-  def __init__(self, channels: List[int], resblk_num: int, nempty: bool = False):
-    super().__init__()
-    self.channels = channels
-    self.resblk_num = resblk_num
-    bottlenecks = [4 if c < 256 else 8 for c in channels]  # to save parameters
-    self.resblocks = torch.nn.ModuleList([
-        ocnn.modules.OctreeResBlocks(ch, ch, resblk_num, bnk, nempty=nempty)
-        for ch, bnk in zip(channels, bottlenecks)])
-
-  def forward(self, datas: List[torch.Tensor], octree: Octree, depth: int):
-    num = len(self.channels)
-    torch._assert(len(datas) == num, 'Error')
-
-    out = [None] * num
-    for i in range(num):
-      depth_i = depth - i
-      out[i] = self.resblocks[i](datas[i], octree, depth_i)
-    return out
-
-
-class TransFunc(torch.nn.Module):
-
-  def __init__(self, in_channels: int, out_channels: int, nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.nempty = nempty
-    self.maxpool = ocnn.nn.OctreeMaxPool(nempty=nempty)
-    self.upsample = ocnn.nn.OctreeUpsample(method='nearest', nempty=nempty)
-    if in_channels != out_channels:
-      self.conv1x1 = ocnn.modules.Conv1x1BnRelu(in_channels, out_channels)
-
-  def forward(self, data: torch.Tensor, octree: Octree,
-              in_depth: int, out_depth: int):
-    out = data
-    if in_depth > out_depth:
-      for d in range(in_depth, out_depth, -1):
-        out = self.maxpool(out, octree, d)
-      if self.in_channels != self.out_channels:
-        out = self.conv1x1(out)
-
-    if in_depth < out_depth:
-      if self.in_channels != self.out_channels:
-        out = self.conv1x1(out)
-      for d in range(in_depth, out_depth, 1):
-        out = self.upsample(out, octree, d)
-    return out
-
-
-class Transitions(torch.nn.Module):
-
-  def __init__(self, channels: List[int], nempty: bool = False):
-    super().__init__()
-    self.channels = channels
-    self.nempty = nempty
-
-    num = len(self.channels)
-    self.trans_func = torch.nn.ModuleList()
-    for i in range(num - 1):
-      for j in range(num):
-        self.trans_func.append(TransFunc(channels[i], channels[j], nempty))
-
-  def forward(self, data: List[torch.Tensor], octree: Octree, depth: int):
-    num = len(self.channels)
-    features = [[None] * (num - 1) for _ in range(num)]
-    for i in range(num - 1):
-      for j in range(num):
-        k = i * num + j
-        in_depth = depth - i
-        out_depth = depth - j
-        features[j][i] = self.trans_func[k](
-            data[i], octree, in_depth, out_depth)
-
-    out = [None] * num
-    for j in range(num):
-      # In the original tensorflow implmentation, a relu is added after the sum.
-      out[j] = torch.stack(features[j], dim=0).sum(dim=0)
-    return out
-
-
-class FrontLayer(torch.nn.Module):
-
-  def __init__(self, channels: List[int], nempty: bool = False):
-    super().__init__()
-    self.channels = channels
-    self.num = len(channels) - 1
-    self.nempty = nempty
-
-    self.conv = torch.nn.ModuleList([
-        ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
-        for i in range(self.num)])
-    self.maxpool = torch.nn.ModuleList([
-        ocnn.nn.OctreeMaxPool(nempty) for i in range(self.num - 1)])
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    out = data
-    for i in range(self.num - 1):
-      depth_i = depth - i
-      out = self.conv[i](out, octree, depth_i)
-      out = self.maxpool[i](out, octree, depth_i)
-    out = self.conv[-1](out, octree, depth - self.num + 1)
-    return out
-
-
-class ClsHeader(torch.nn.Module):
-
-  def __init__(self, channels: List[int], out_channels: int, nempty: bool = False):
-    super().__init__()
-    self.channels = channels
-    self.out_channels = out_channels
-    self.nempty = nempty
-
-    in_channels = int(torch.Tensor(channels).sum())
-    self.conv1x1 = ocnn.modules.Conv1x1BnRelu(in_channels, 1024)
-    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
-    self.header = torch.nn.Sequential(
-        torch.nn.Flatten(start_dim=1),
-        torch.nn.Linear(1024, out_channels, bias=True))
-    # self.header = torch.nn.Sequential(
-    #     ocnn.modules.FcBnRelu(512, 256),
-    #     torch.nn.Dropout(p=0.5),
-    #     torch.nn.Linear(256, out_channels))
-
-  def forward(self, data: List[torch.Tensor], octree: Octree, depth: int):
-    full_depth = 2
-    num = len(data)
-    outs = [x for x in data]  # avoid modifying the input data
-    for i in range(num):
-      depth_i = depth - i
-      for d in range(depth_i, full_depth, -1):
-        outs[i] = ocnn.nn.octree_max_pool(outs[i], octree, d, self.nempty)
-
-    out = torch.cat(outs, dim=1)
-    out = self.conv1x1(out)
-    out = self.global_pool(out, octree, full_depth)
-    logit = self.header(out)
-    return logit
-
-
-class HRNet(torch.nn.Module):
-  r''' Octree-based HRNet for classification and segmentation. '''
-
-  def __init__(self, in_channels: int, out_channels: int, stages: int = 3,
-               interp: str = 'linear', nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.interp = interp
-    self.nempty = nempty
-    self.stages = stages
-
-    self.resblk_num = 3
-    self.channels = [128, 256, 512, 512]
-
-    self.front = FrontLayer([in_channels, 32, self.channels[0]], nempty)
-    self.branches = torch.nn.ModuleList([
-        Branches(self.channels[:i+1], self.resblk_num, nempty)
-        for i in range(stages)])
-    self.transitions = torch.nn.ModuleList([
-        Transitions(self.channels[:i+2], nempty)
-        for i in range(stages-1)])
-
-    self.cls_header = ClsHeader(self.channels[:stages], out_channels, nempty)
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-    convs = [self.front(data, octree, depth)]
-    depth = depth - 1  # the data is downsampled in `front`
-    for i in range(self.stages):
-      convs = self.branches[i](convs, octree, depth)
-      if i < self.stages - 1:
-        convs = self.transitions[i](convs, octree, depth)
-
-    logits = self.cls_header(convs, octree, depth)
-
-    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
+from typing import List
+
+import ocnn
+from ocnn.octree import Octree
+
+
+class Branches(torch.nn.Module):
+
+  def __init__(self, channels: List[int], resblk_num: int, nempty: bool = False):
+    super().__init__()
+    self.channels = channels
+    self.resblk_num = resblk_num
+    bottlenecks = [4 if c < 256 else 8 for c in channels]  # to save parameters
+    self.resblocks = torch.nn.ModuleList([
+        ocnn.modules.OctreeResBlocks(ch, ch, resblk_num, bnk, nempty=nempty)
+        for ch, bnk in zip(channels, bottlenecks)])
+
+  def forward(self, datas: List[torch.Tensor], octree: Octree, depth: int):
+    num = len(self.channels)
+    torch._assert(len(datas) == num, 'Error')
+
+    out = [None] * num
+    for i in range(num):
+      depth_i = depth - i
+      out[i] = self.resblocks[i](datas[i], octree, depth_i)
+    return out
+
+
+class TransFunc(torch.nn.Module):
+
+  def __init__(self, in_channels: int, out_channels: int, nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.nempty = nempty
+    self.maxpool = ocnn.nn.OctreeMaxPool(nempty=nempty)
+    self.upsample = ocnn.nn.OctreeUpsample(method='nearest', nempty=nempty)
+    if in_channels != out_channels:
+      self.conv1x1 = ocnn.modules.Conv1x1BnRelu(in_channels, out_channels)
+
+  def forward(self, data: torch.Tensor, octree: Octree,
+              in_depth: int, out_depth: int):
+    out = data
+    if in_depth > out_depth:
+      for d in range(in_depth, out_depth, -1):
+        out = self.maxpool(out, octree, d)
+      if self.in_channels != self.out_channels:
+        out = self.conv1x1(out)
+
+    if in_depth < out_depth:
+      if self.in_channels != self.out_channels:
+        out = self.conv1x1(out)
+      for d in range(in_depth, out_depth, 1):
+        out = self.upsample(out, octree, d)
+    return out
+
+
+class Transitions(torch.nn.Module):
+
+  def __init__(self, channels: List[int], nempty: bool = False):
+    super().__init__()
+    self.channels = channels
+    self.nempty = nempty
+
+    num = len(self.channels)
+    self.trans_func = torch.nn.ModuleList()
+    for i in range(num - 1):
+      for j in range(num):
+        self.trans_func.append(TransFunc(channels[i], channels[j], nempty))
+
+  def forward(self, data: List[torch.Tensor], octree: Octree, depth: int):
+    num = len(self.channels)
+    features = [[None] * (num - 1) for _ in range(num)]
+    for i in range(num - 1):
+      for j in range(num):
+        k = i * num + j
+        in_depth = depth - i
+        out_depth = depth - j
+        features[j][i] = self.trans_func[k](
+            data[i], octree, in_depth, out_depth)
+
+    out = [None] * num
+    for j in range(num):
+      # In the original tensorflow implmentation, a relu is added after the sum.
+      out[j] = torch.stack(features[j], dim=0).sum(dim=0)
+    return out
+
+
+class FrontLayer(torch.nn.Module):
+
+  def __init__(self, channels: List[int], nempty: bool = False):
+    super().__init__()
+    self.channels = channels
+    self.num = len(channels) - 1
+    self.nempty = nempty
+
+    self.conv = torch.nn.ModuleList([
+        ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
+        for i in range(self.num)])
+    self.maxpool = torch.nn.ModuleList([
+        ocnn.nn.OctreeMaxPool(nempty) for i in range(self.num - 1)])
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    out = data
+    for i in range(self.num - 1):
+      depth_i = depth - i
+      out = self.conv[i](out, octree, depth_i)
+      out = self.maxpool[i](out, octree, depth_i)
+    out = self.conv[-1](out, octree, depth - self.num + 1)
+    return out
+
+
+class ClsHeader(torch.nn.Module):
+
+  def __init__(self, channels: List[int], out_channels: int, nempty: bool = False):
+    super().__init__()
+    self.channels = channels
+    self.out_channels = out_channels
+    self.nempty = nempty
+
+    in_channels = int(torch.Tensor(channels).sum())
+    self.conv1x1 = ocnn.modules.Conv1x1BnRelu(in_channels, 1024)
+    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
+    self.header = torch.nn.Sequential(
+        torch.nn.Flatten(start_dim=1),
+        torch.nn.Linear(1024, out_channels, bias=True))
+    # self.header = torch.nn.Sequential(
+    #     ocnn.modules.FcBnRelu(512, 256),
+    #     torch.nn.Dropout(p=0.5),
+    #     torch.nn.Linear(256, out_channels))
+
+  def forward(self, data: List[torch.Tensor], octree: Octree, depth: int):
+    full_depth = 2
+    num = len(data)
+    outs = [x for x in data]  # avoid modifying the input data
+    for i in range(num):
+      depth_i = depth - i
+      for d in range(depth_i, full_depth, -1):
+        outs[i] = ocnn.nn.octree_max_pool(outs[i], octree, d, self.nempty)
+
+    out = torch.cat(outs, dim=1)
+    out = self.conv1x1(out)
+    out = self.global_pool(out, octree, full_depth)
+    logit = self.header(out)
+    return logit
+
+
+class HRNet(torch.nn.Module):
+  r''' Octree-based HRNet for classification and segmentation. '''
+
+  def __init__(self, in_channels: int, out_channels: int, stages: int = 3,
+               interp: str = 'linear', nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.interp = interp
+    self.nempty = nempty
+    self.stages = stages
+
+    self.resblk_num = 3
+    self.channels = [128, 256, 512, 512]
+
+    self.front = FrontLayer([in_channels, 32, self.channels[0]], nempty)
+    self.branches = torch.nn.ModuleList([
+        Branches(self.channels[:i+1], self.resblk_num, nempty)
+        for i in range(stages)])
+    self.transitions = torch.nn.ModuleList([
+        Transitions(self.channels[:i+2], nempty)
+        for i in range(stages-1)])
+
+    self.cls_header = ClsHeader(self.channels[:stages], out_channels, nempty)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+    convs = [self.front(data, octree, depth)]
+    depth = depth - 1  # the data is downsampled in `front`
+    for i in range(self.stages):
+      convs = self.branches[i](convs, octree, depth)
+      if i < self.stages - 1:
+        convs = self.transitions[i](convs, octree, depth)
+
+    logits = self.cls_header(convs, octree, depth)
+
+    return logits

ocnn/models/image2shape.py

@@ -0,0 +1,128 @@
+# --------------------------------------------------------
+# 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 torch.nn
+from typing import Optional
+from torchvision.models import resnet18
+
+import ocnn
+from ocnn.octree import Octree
+
+
+class Image2Shape(torch.nn.Module):
+  r''' Octree-based AutoEncoder for shape encoding and decoding.
+
+  Args:
+    channel_out (int): The channel of the output signal.
+    depth (int): The depth of the octree.
+    full_depth (int): The full depth of the octree.
+  '''
+
+  def __init__(self, channel_out: int, depth: int, full_depth: int = 2,
+               code_channel: int = 32):
+    super().__init__()
+    self.depth = depth
+    self.full_depth = full_depth
+    self.channel_out = channel_out
+    self.resblk_num = 2
+    self.channels = [512, 512, 256, 256, 128, 128, 64, 64, 32, 32]
+    self.code_channel = code_channel
+
+    # encoder
+    self.resnet18 = resnet18()
+    channel = self.code_channel * 2 ** (3 * full_depth)
+    self.resnet18.fc = torch.nn.Linear(512, channel, bias=True)
+
+    # decoder
+    self.channels[full_depth] = self.code_channel  # update `channels`
+    self.upsample = torch.nn.ModuleList([ocnn.modules.OctreeDeconvBnRelu(
+        self.channels[d-1], self.channels[d], kernel_size=[2], stride=2,
+        nempty=False) for d in range(full_depth+1, depth+1)])
+    self.decoder_blks = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
+        self.channels[d], self.channels[d], self.resblk_num, nempty=False)
+        for d in range(full_depth, depth+1)])
+
+    # header
+    self.predict = torch.nn.ModuleList([self._make_predict_module(
+        self.channels[d], 2) for d in range(full_depth, depth + 1)])
+    self.header = self._make_predict_module(self.channels[depth], channel_out)
+
+  def _make_predict_module(self, channel_in, channel_out=2, num_hidden=64):
+    return torch.nn.Sequential(
+        ocnn.modules.Conv1x1BnRelu(channel_in, num_hidden),
+        ocnn.modules.Conv1x1(num_hidden, channel_out, use_bias=True))
+
+  def decoder(self, shape_code: torch.Tensor, octree: Octree,
+              update_octree: bool = False):
+    r''' The decoder network of the AutoEncoder.
+    '''
+
+    logits = dict()
+    deconv = shape_code
+    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, d-1)
+      deconv = self.decoder_blks[i](deconv, octree, 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.octree_split(split, d)
+        if d < depth:
+          octree.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, depth)
+        if update_octree:
+          octree.features[depth] = signal
+
+    return {'logits': logits, 'signal': signal, 'octree_out': octree}
+
+  def decode_code(self, shape_code: torch.Tensor):
+    r''' Decodes the shape code to an output octree.
+
+    Args:
+      shape_code (torch.Tensor): The shape code for decoding.
+    '''
+
+    octree_out = self.init_octree(shape_code)
+    out = self.decoder(shape_code, octree_out, update_octree=True)
+    return out
+
+  def init_octree(self, shape_code: torch.Tensor):
+    r''' Initialize a full octree for decoding.
+
+    Args:
+      shape_code (torch.Tensor): The shape code for decoding, used to getting
+          the `batch_size` and `device` to initialize the output octree.
+    '''
+
+    node_num = 2 ** (3 * self.full_depth)
+    batch_size = shape_code.size(0) // node_num
+    octree = ocnn.octree.init_octree(
+        self.depth, self.full_depth, batch_size, shape_code.device)
+    return octree
+
+  def forward(self, image: torch.Tensor, octree: Optional[Octree] = None,
+              update_octree: bool = False):
+    r''''''
+
+    shape_code = self.resnet18(image)
+    shape_code = shape_code.view(-1, self.code_channel)
+    if update_octree:
+      octree = self.init_octree(shape_code)
+    out = self.decoder(shape_code, octree, update_octree)
+    return out

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 i 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 i 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