openpoints 0.1.0__py3-none-any.whl

openpoints/models/reconstruction/base_recontruct.py

@@ -0,0 +1,216 @@
+""" Maksed PointViT in PyTorch
+Copyright 2022@PointNeXt team
+"""
+import logging
+import torch
+import torch.nn as nn
+from ..build import MODELS
+from ..layers.attention import Block
+from ..layers import create_norm, create_linearblock
+from ..layers.graph_conv import GraphConv, DilatedKNN
+
+
+@MODELS.register_module()
+class MaskedTransformerDecoder(nn.Module):
+    """ MaskedTransformerDecoder
+    """
+
+    def __init__(self,
+                 embed_dim,  # the out features dim of the encoder
+                 group_size=32,  # the number of points inside each group
+                 decoder_embed_dim=192, decoder_depth=4, decoder_num_heads=16,
+                 norm_args={'norm': 'ln', 'eps': 1.0e-6},
+                 act_args={'act': 'gelu'},
+                 add_pos_each_block=True, 
+                 **kwargs
+                 ):
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+
+        # ------------------------------------------------------------------
+        # MAE decoder
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+        self.mask_token = nn.Parameter(torch.randn(1, 1, decoder_embed_dim))
+        self.decoder_cls_pos = nn.Parameter(torch.randn(1, 1, decoder_embed_dim))
+        self.decoder_pos_embed = nn.Sequential(
+            nn.Linear(3, 128),
+            nn.GELU(),
+            nn.Linear(128, decoder_embed_dim)
+        )
+        self.add_pos_each_block = add_pos_each_block
+        self.decoder_blocks = nn.ModuleList([
+            Block(
+                dim=decoder_embed_dim, num_heads=decoder_num_heads,
+                norm_args=norm_args, act_args=act_args
+            )
+            for _ in range(decoder_depth)])
+        self.decoder_norm = create_norm(norm_args, decoder_embed_dim)
+        self.decoder_pred = nn.Linear(decoder_embed_dim, group_size * 3, bias=True)  # decoder to patch
+        # ------------------------------------------------------------------
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        torch.nn.init.normal_(self.decoder_cls_pos, std=.02)
+        torch.nn.init.normal_(self.mask_token, std=.02)
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d, nn.BatchNorm1d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, center_xyz, features, ids_restore):
+        # embed tokens
+        features = self.decoder_embed(features)
+        B, L, C = features.shape  # batch size, length, channels
+        
+        # in the decoder part. we know the positional encoding of groups
+        decoder_pos_embed = torch.cat(
+            (self.decoder_cls_pos.expand(B, -1, -1), self.decoder_pos_embed(center_xyz)), dim=1)
+
+        # append mask tokens to sequence
+        # use mask tokens to fill the masked features. this is why missing part should be 1 not 0.
+        mask_tokens = self.mask_token.repeat(B, ids_restore.shape[1] + 1 - L, 1)    # +1, since features contains additional cls token. 
+        x_ = torch.cat([features[:, 1:, :], mask_tokens], dim=1)  # no cls token
+        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, C))  # unshuffle
+        features = torch.cat([features[:, :1, :], x_], dim=1)  # append cls token
+
+        if self.add_pos_each_block:
+            for block in self.decoder_blocks:
+                features = block(features + decoder_pos_embed)
+        else:
+            features = self.pos_drop(features + decoder_pos_embed)
+            for block in self.decoder_blocks:
+                features = block(features)
+
+        features = self.decoder_norm(features)
+        # predictor projection
+        features = self.decoder_pred(features)
+        # remove cls token
+        features = features[:, 1:, :]
+        return features
+
+
+@MODELS.register_module()
+class FoldingNet(nn.Module):
+    """ FoldingNet.
+    Used in many methods, e.g. FoldingNet, PCN, OcCo, Point-BERT
+    learning point reconstruction only from global feature
+    """
+    def __init__(self, in_channels, emb_dims=1024,
+                 num_fine=1024,
+                 grid_size=2,
+                 **kwargs):
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+
+        self.grid_size = grid_size
+        self.num_coarse = num_fine // grid_size**2
+        self.num_fine = num_fine
+
+        self.folding1 = nn.Sequential(
+            nn.Linear(in_channels, emb_dims),
+            nn.ReLU(inplace=True),
+            nn.Linear(emb_dims, emb_dims),
+            nn.ReLU(inplace=True),
+            nn.Linear(emb_dims, self.num_coarse * 3))
+
+        self.folding2 = nn.Sequential(
+            nn.Linear(emb_dims+2+3, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 3))
+
+        a = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(1, self.grid_size).expand(self.grid_size, self.grid_size).reshape(1, -1)
+        b = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(self.grid_size, 1).expand(self.grid_size, self.grid_size).reshape(1, -1)
+        self.register_buffer('folding_seed', torch.cat([a, b], dim=0).reshape(1, 2, self.grid_size ** 2).transpose(1, 2))
+        self.model_init()
+
+    def model_init(self):
+        for m in self.modules():
+            if isinstance(m, (torch.nn.Conv2d, torch.nn.Conv1d)):
+                torch.nn.init.kaiming_normal_(m.weight)
+                m.weight.requires_grad = True
+                if m.bias is not None:
+                    m.bias.data.zero_()
+                    m.bias.requires_grad = True
+            elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d, nn.BatchNorm1d)):
+                nn.init.constant_(m.bias, 0)
+                nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, xyz, x, **kwargs):
+        B = x.shape[0]
+        coarse = self.folding1(x)
+        coarse = coarse.view(-1, self.num_coarse, 3)
+        point_feat = coarse.unsqueeze(2).expand(-1, -1, self.grid_size**2, -1).reshape([-1, self.num_fine, 3])
+
+        seed = self.folding_seed.unsqueeze(1).expand(B, self.num_coarse, -1, -1).reshape(B, self.num_fine, -1)
+
+        x = x.unsqueeze(1).expand(-1, self.num_fine, -1)
+        feat = torch.cat([x, seed, point_feat], dim=-1)
+
+        center = coarse.unsqueeze(2).expand(-1, -1, self.grid_size**2, -1).reshape([-1, self.num_fine, 3])
+
+        fine = self.folding2(feat) + center
+
+        return coarse, fine
+
+
+@MODELS.register_module()
+class NodeShuffle(nn.Module):
+    """ NodeShuffle
+        proposed in PU-GCN
+    """
+
+    def __init__(self,
+                 in_channels,
+                 up_ratio=16,
+                 emb_dims=1024,
+                 k=16,
+                 norm_args={'norm': 'bn'},
+                 act_args={'act': 'relu'},
+                 **kwargs):
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+
+        self.up_ratio = up_ratio
+        conv = 'edge'
+        self.knn = DilatedKNN(k, 1)
+        self.convs = nn.Sequential(
+            GraphConv(in_channels, emb_dims, conv, norm_args=norm_args, act_args=act_args),
+            GraphConv(emb_dims, emb_dims, conv, norm_args=norm_args, act_args=act_args)
+        )
+        self.proj = create_linearblock(emb_dims, 3 * up_ratio)
+        self.model_init()
+
+    def model_init(self):
+        for m in self.modules():
+            if isinstance(m, (torch.nn.Conv2d, torch.nn.Conv1d)):
+                torch.nn.init.kaiming_normal_(m.weight)
+                m.weight.requires_grad = True
+                if m.bias is not None:
+                    m.bias.data.zero_()
+                    m.bias.requires_grad = True
+            elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d, nn.BatchNorm1d)):
+                nn.init.constant_(m.bias, 0)
+                nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, xyz, feature, **kwargs):
+        # learn displacement
+        B, C, N = feature.shape
+        feature = feature.unsqueeze(-1)
+        edge_index = self.knn(xyz)
+        for conv in self.convs:
+            feature = conv(feature, edge_index)
+        new_xyz = self.proj(feature.squeeze(-1).transpose(1, 2)).view(B, N, -1, 3) + xyz.unsqueeze(2).repeat(-1, -1, self.up_ratio, -1)
+        return new_xyz.view(B, -1, 3)

openpoints/models/reconstruction/maskedpoint.py

@@ -0,0 +1,116 @@
+""" Maksed PointViT in PyTorch
+Copyright 2022@PointNeXt team
+"""
+import logging
+import torch
+import torch.nn as nn
+from openpoints.cpp.chamfer_dist import ChamferDistanceL1
+from ..build import build_model_from_cfg, MODELS
+
+
+@MODELS.register_module()
+class MaskedPoint(nn.Module):
+    """ Masked AutoEncoder for Point-based methods
+    """
+    def __init__(self,
+                 backbone_args,
+                 decoder_args,
+                 mask_ratio,
+                 **kwargs
+                 ):
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+
+        # ------------------------------------------------------------------
+        # MAE encoder. e.g. DGCNN, DeepGCN, PointNet++
+        self.encoder = build_model_from_cfg(backbone_args)
+
+        # ------------------------------------------------------------------
+        # MAE decoder. e.g. FoldingNet (works bad in random sampling), PU-GCN (NodeShuffle)
+        self.use_global_feat = True if decoder_args.NAME.lower() in ['foldingnet', 'pointcompletion'] else False
+        self.maxpool = lambda x: torch.max(x, dim=-1, keepdim=False)[0]
+        self.decoder = build_model_from_cfg(decoder_args)
+
+        # ------------------------------------------------------------------
+        # loss
+        self.mask_ratio = mask_ratio
+        self.build_loss_func()
+
+    @staticmethod
+    def random_masking(x, features=None, mask_ratio=0.9):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, 3], sequence
+        features: [N, D, L], sequence
+        TODO: suppport other masking. Like OcCo, block masking as ablation
+        """
+        N, L, _ = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).expand(-1, -1, 3))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        # features
+        if features is not None:
+            features_masked = torch.gather(x, dim=2, index=ids_keep.unsqueeze(1).expand(-1, features.shape[1], -1))
+        else:
+            features_masked = None
+        return x_masked, features_masked, mask, ids_restore, ids_keep
+
+    def build_loss_func(self, smoothing=False):
+        self.criterion = ChamferDistanceL1()
+
+    def forward_loss(self, xyz, pred, mask=None):
+        """
+        # chamfer distance. two options 
+        1. chamfer distance on the merged point clouds. 
+        2. chamfer distance per local patch
+        
+        
+        xyz: the original points [B, N, 3] 
+        grouped_xyz: the points after grouping. [B, 3, L, K]
+        pred: [B, L, K*3]
+        mask: [B, L], 0 is keep, 1 is remove,
+        idx_keep: [B, L]
+        """
+        if isinstance(pred, (tuple, list)):
+            loss = 0.
+            for pred_i in pred:
+                loss += self.criterion(pred_i, xyz)
+        else:
+            loss = self.criterion(pred, xyz)
+        # mask = mask.unsqueeze(-1)
+        # grouped_xyz = torch.mul(grouped_xyz.permute(0, 2, 3, 1), mask.unsqueeze(-1)).reshape(-1, K, C)    
+        # pred = torch.mul(pred, mask).reshape(-1, K, C)  
+        # loss = self.criterion(pred, grouped_xyz)/ mask.sum()
+        return loss
+
+    def forward(self, xyz, features=None):
+        xyz_masked, features_masked, mask, ids_restore, idx_keep = self.random_masking(xyz, features, self.mask_ratio)
+        latent = self.encoder(xyz_masked, features_masked)
+        if self.use_global_feat:
+            latent = self.maxpool(latent)
+        pred = self.decoder(xyz_masked, latent, ids_restore)  # [N, L, p*p*3]
+
+        """visualize pred. TODO: add to Tensorboard 
+            from openpoints.dataset import vis_multi_points
+            vis_multi_points((xyz[0].cpu().detach().numpy(), xyz_masked[0].cpu().detach().numpy(), pred[-1][0].cpu().detach().numpy()))
+        """
+        loss = self.forward_loss(xyz, pred, mask)
+        return loss, pred
+

openpoints/models/reconstruction/maskedpointgroup.py

@@ -0,0 +1,168 @@
+""" Maksed PointViT in PyTorch
+Copyright 2022@PointNeXt team
+"""
+import logging
+import torch
+import torch.nn as nn
+from openpoints.cpp.chamfer_dist import ChamferDistanceL1
+from ..build import build_model_from_cfg, MODELS
+from ..layers.subsample import furthest_point_sample, random_sample
+from ..layers.group import KNNGroup, QueryAndGroup
+
+
+@MODELS.register_module()
+class MaskedPointGroup(nn.Module):
+    """ Masked AutoEncoder for Point-based methods
+    """
+
+    def __init__(self,
+                 encoder_args,
+                 decoder_args,
+                 mask_ratio,
+                 subsample='fps',  # random, FPS
+                 group='knn', 
+                 group_size=32,
+                 sample_ratio=0.25,
+                 radius=0.1,
+                 add_cls_token=False,
+                 **kwargs
+                 ):
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+
+        # ------------------------------------------------------------------
+        # Grouping
+        self.group_size = group_size
+        self.sample_ratio = sample_ratio  # downsample 4x
+        if subsample.lower() == 'fps':
+            self.sample_fn = furthest_point_sample
+        elif 'random' in subsample.lower():
+            self.sample_fn = random_sample
+
+        self.group = group.lower()
+        if 'ball' in self.group or 'query' in self.group:
+            self.grouper = QueryAndGroup(nsample=self.group_size,
+                                         relative_p=False, normalize_p=False,
+                                         radius=radius)
+        elif 'knn' in self.group.lower():
+            self.grouper = KNNGroup(self.group_size, relative_p=False)
+        else:
+            raise NotImplementedError(f'{self.group.lower()} is not implemented. Only support ballquery, knn')
+
+        # ------------------------------------------------------------------
+        # MAE encoder. e.g. DGCNN, DeepGCN, PointNet++
+        self.encoder = build_model_from_cfg(encoder_args)
+        self.add_cls_token = add_cls_token
+        if self.add_cls_token:
+            self.cls_token = nn.Parameter(torch.randn(1, 1, decoder_args.embed_dim))
+            torch.nn.init.normal_(self.cls_token, std=.02)
+        # ------------------------------------------------------------------
+        # MAE decoder. e.g. FoldingNet (works bad in random sampling), PU-GCN (NodeShuffle)
+        self.use_global_feat = True if decoder_args.NAME.lower() in ['foldingnet', 'pointcompletion'] else False
+        self.maxpool = lambda x: torch.max(x, dim=-1, keepdim=False)[0]
+        self.decoder = build_model_from_cfg(decoder_args)
+
+        # ------------------------------------------------------------------
+        # loss
+        self.mask_ratio = mask_ratio
+        self.build_loss_func()
+
+    @staticmethod
+    def group_random_masking(x, f=None, mask_ratio=0.9):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, 3], sequence
+        f: [N, D, L], sequence
+        """
+        B, _, N, K = x.shape  # batch, dim, num_points, num_neighbors
+        len_keep = int(N * (1 - mask_ratio))
+
+        noise = torch.rand(B, N, device=x.device)  # noise in [0, 1]
+
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=2, index=ids_keep.unsqueeze(1).unsqueeze(-1).expand(-1, 3, -1, K))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([B, N], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        # f
+        if f is not None:
+            f_masked = torch.gather(f, dim=2,
+                                           index=ids_keep.unsqueeze(1).unsqueeze(-1).expand(-1, f.shape[1], -1,
+                                                                                            f.shape[-1]))
+        else:
+            f_masked = None
+        return x_masked, f_masked, mask, ids_restore, ids_keep
+
+    def build_loss_func(self, smoothing=False):
+        self.criterion = ChamferDistanceL1()
+
+    def forward_loss(self, dp, pred, mask=None):
+        """
+        # chamfer distance. two options
+        1. chamfer distance on the merged point clouds.
+        2. chamfer distance per local patch
+
+
+        p: the original points [B, N, 3]
+        dp: the points after grouping. [B, 3, L, K]
+        pred: [B, L, K*3]
+        mask: [B, L], 0 is keep, 1 is remove,
+        idx_keep: [B, L]
+        """
+        # option 2, per patch chamfer distance.
+        B, C, L, K = dp.shape
+        # reshape dp and pred as [BL, K, C]
+        dp = dp.permute(0, 2, 3, 1).reshape(-1, K, C)
+        pred = pred.reshape(-1, K, C)
+        loss = self.criterion(pred, dp)
+        return loss
+
+    def forward(self, p, f=None):
+        # downsample, N -> L, e.g. 1024 -> 256
+        if isinstance(p, dict):
+            p, f = p['pos'], p.get('x', None)
+        if f is None:
+            f = p.transpose(1, 2).contiguous()
+        B, N, _ = p.shape[:3]
+        idx = self.sample_fn(p, int(N * self.sample_ratio)).long()
+        center_p = torch.gather(p, 1, idx.unsqueeze(-1).expand(-1, -1, 3))  # e.g.  [B, L, 3]
+
+        # query neighbors. dp: [B, 3, L, K], a typical K is 8
+        dp, gf = self.grouper(center_p, p, f)
+        dp_masked, gf_masked, mask, ids_restore, idx_keep = self.group_random_masking(
+            dp, gf, self.mask_ratio)
+
+        gf_masked = torch.cat((dp_masked, gf_masked), dim=1)
+        latent = self.encoder.ssl_forward(dp_masked, gf_masked)  # latent: [B, C, MK]
+        if self.use_global_feat:
+            latent = self.maxpool(latent)
+        else:
+            latent = latent.transpose(1, 2)
+
+        # concat token
+        if self.add_cls_token:
+            latent = torch.cat((self.cls_token.expand(B, -1, -1), latent), dim=1)
+        pred = self.decoder(center_p, latent, ids_restore)  # [N, L, p*p*3]
+
+        # pred is the reconsructed grouped p.
+        """visualize pred. TODO: add to Tensorboard 
+            from openpoints.dataset import vis_multi_points
+            B, C, L, K = dp.shape
+            input_group_pts = dp.permute(0, 2, 3, 1).reshape(B, -1, 3).detach().cpu().numpy()
+            pred_group_pts = pred.clone().reshape(B, -1, 3).detach().cpu().numpy()
+            # vis_multi_points((p[0].cpu().numpy(), center_p[0].detach().cpu().numpy(), dp_masked[0].detach().cpu().numpy(), pred_group_pts[0]))
+            vis_multi_points((center_p[0].detach().cpu().numpy(), input_group_pts[0], dp_masked[0].detach().cpu().numpy(), pred_group_pts[0]))
+        """
+        loss = self.forward_loss(dp, pred, mask)
+        return loss, pred