openpoints 0.1.0__py3-none-any.whl

openpoints/models/backbone/simpleview.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn as nn
+from .simpleview_util import PCViews
+from ..build import MODELS
+# from roi_align import CropAndResize # crop_and_resize module
+from openpoints.transforms import build_transforms_from_cfg
+
+
+class Squeeze(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, inp):
+        return inp.squeeze()
+
+
+class BatchNormPoint(nn.Module):
+    def __init__(self, feat_size):
+        super().__init__()
+        self.feat_size = feat_size
+        self.bn = nn.BatchNorm1d(feat_size)
+
+    def forward(self, x):
+        assert len(x.shape) == 3
+        s1, s2, s3 = x.shape[0], x.shape[1], x.shape[2]
+        assert s3 == self.feat_size
+        x = x.view(s1 * s2, self.feat_size)
+        x = self.bn(x)
+        return x.view(s1, s2, s3)
+
+
+@MODELS.register_module()
+class MVFC(nn.Module):
+    """
+    Final FC layers for the MV model
+    """
+
+    def __init__(self, num_views, in_features, out_features, dropout):
+        super().__init__()
+        self.num_views = num_views
+        self.in_features = in_features
+        self.model = nn.Sequential(
+            BatchNormPoint(in_features),
+            # dropout before concatenation so that each view drops features independently
+            nn.Dropout(dropout),
+            nn.Flatten(),
+            nn.Linear(in_features=in_features * self.num_views,
+                      out_features=in_features),
+            nn.BatchNorm1d(in_features),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(in_features=in_features, out_features=out_features,
+                      bias=True))
+
+    def forward(self, feat):
+        feat = feat.view((-1, self.num_views, self.in_features))
+        out = self.model(feat)
+        return out
+
+
+@MODELS.register_module()
+class MVModel(nn.Module):
+    def __init__(self, task='cls', backbone='resnet18',
+                 channels=16,
+                 num_classes=15,
+                 resolution=128,
+                 use_img_transform=False,
+                 **kwargs):
+        super().__init__()
+        assert task == 'cls'
+        self.task = task
+        self.num_classes = num_classes
+        self.dropout = kwargs.get('dropout', 0.5)
+        self.channels = channels
+        pc_views = PCViews()
+        self.num_views = pc_views.num_views
+        self._get_img = pc_views.get_img
+
+        img_layers, in_features = self.get_img_layers(
+            backbone, channels=channels)
+        self.img_model = nn.Sequential(*img_layers)
+
+        self.final_fc = MVFC(
+            num_views=self.num_views,
+            in_features=in_features,
+            out_features=self.num_classes,
+            dropout=self.dropout)
+        if use_img_transform:
+            self.img_transform = build_transforms_from_cfg('img', {'img': ['Zoom']})
+        else:
+            self.img_transform = None
+            
+    def forward(self, pc):
+        """
+        :param pc:
+        :return:
+        """
+        if hasattr(pc, 'keys'):
+            pc = pc['pos']
+        img = self.get_img(pc)
+
+        if self.training and self.img_transform is not None:
+            img = self.img_transform(img, pc.shape[0] * self.num_views)
+        feat = self.img_model(img)
+        logit = self.final_fc(feat)
+        return logit
+
+    def forward_cls_feat(self, pc):
+        return self.forward(pc)
+
+    def get_img(self, pc):
+        img = self._get_img(pc)
+        img = img.unsqueeze(3)
+        # [num_pc * num_views, 1, RESOLUTION, RESOLUTION]
+        img = img.permute(0, 3, 1, 2)
+        return img
+
+    @staticmethod
+    def get_img_layers(backbone, channels):
+        """
+        Return layers for the image model
+        """
+
+        from .resnet import _resnet, BasicBlock
+        assert backbone == 'resnet18'
+        layers = [2, 2, 2, 2]
+        block = BasicBlock
+        backbone_mod = _resnet(
+            arch=None,
+            block=block,
+            layers=layers,
+            pretrained=False,
+            progress=False,
+            feature_size=channels,
+            zero_init_residual=True)
+
+        all_layers = [x for x in backbone_mod.children()]
+        in_features = all_layers[-1].in_features
+
+        # all layers except the final fc layer and the initial conv layers
+        # WARNING: this is checked only for resnet models
+        main_layers = all_layers[4:-1]
+        img_layers = [
+            nn.Conv2d(1, channels, kernel_size=(3, 3), stride=(1, 1),
+                      padding=(1, 1), bias=False),
+            nn.BatchNorm2d(channels, eps=1e-05, momentum=0.1,
+                           affine=True, track_running_stats=True),
+            nn.ReLU(inplace=True),
+            *main_layers,
+            Squeeze()
+        ]
+
+        return img_layers, in_features

openpoints/models/backbone/simpleview_util.py

@@ -0,0 +1,292 @@
+import torch.nn as nn
+import numpy as np
+import torch
+
+RESOLUTION = 128
+TRANS = -1.4
+
+def euler2mat(angle):
+    """Convert euler angles to rotation matrix.
+     :param angle: [3] or [b, 3]
+     :return
+        rotmat: [3] or [b, 3, 3]
+    source
+    https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py
+    """
+
+    if len(angle.size()) == 1:
+        x, y, z = angle[0], angle[1], angle[2]
+        _dim = 0
+        _view = [3, 3]
+    elif len(angle.size()) == 2:
+        b, _ = angle.size()
+        x, y, z = angle[:, 0], angle[:, 1], angle[:, 2]
+        _dim = 1
+        _view = [b, 3, 3]
+
+    else:
+        assert False
+
+    cosz = torch.cos(z)
+    sinz = torch.sin(z)
+
+    # zero = torch.zeros([b], requires_grad=False, device=angle.device)[0]
+    # one = torch.ones([b], requires_grad=False, device=angle.device)[0]
+    zero = z.detach()*0
+    one = zero.detach()+1
+    zmat = torch.stack([cosz, -sinz, zero,
+                        sinz, cosz, zero,
+                        zero, zero, one], dim=_dim).reshape(_view)
+
+    cosy = torch.cos(y)
+    siny = torch.sin(y)
+
+    ymat = torch.stack([cosy, zero, siny,
+                        zero, one, zero,
+                        -siny, zero, cosy], dim=_dim).reshape(_view)
+
+    cosx = torch.cos(x)
+    sinx = torch.sin(x)
+
+    xmat = torch.stack([one, zero, zero,
+                        zero, cosx, -sinx,
+                        zero, sinx, cosx], dim=_dim).reshape(_view)
+
+    rot_mat = xmat @ ymat @ zmat
+    # print(rot_mat)
+    return rot_mat
+
+
+def distribute(depth, _x, _y, size_x, size_y, image_height, image_width):
+    """
+    Distributes the depth associated with each point to the discrete coordinates (image_height, image_width) in a region
+    of size (size_x, size_y).
+    :param depth:
+    :param _x:
+    :param _y:
+    :param size_x:
+    :param size_y:
+    :param image_height:
+    :param image_width:
+    :return:
+    """
+
+    assert size_x % 2 == 0 or size_x == 1
+    assert size_y % 2 == 0 or size_y == 1
+    batch, _ = depth.size()
+    epsilon = torch.tensor([1e-12], requires_grad=False, device=depth.device)
+    _i = torch.linspace(-size_x / 2, (size_x / 2) - 1, size_x, requires_grad=False, device=depth.device)
+    _j = torch.linspace(-size_y / 2, (size_y / 2) - 1, size_y, requires_grad=False, device=depth.device)
+
+    extended_x = _x.unsqueeze(2).repeat([1, 1, size_x]) + _i  # [batch, num_points, size_x]
+    extended_y = _y.unsqueeze(2).repeat([1, 1, size_y]) + _j  # [batch, num_points, size_y]
+
+    extended_x = extended_x.unsqueeze(3).repeat([1, 1, 1, size_y])  # [batch, num_points, size_x, size_y]
+    extended_y = extended_y.unsqueeze(2).repeat([1, 1, size_x, 1])  # [batch, num_points, size_x, size_y]
+
+    extended_x.ceil_()
+    extended_y.ceil_()
+
+    value = depth.unsqueeze(2).unsqueeze(3).repeat([1, 1, size_x, size_y])  # [batch, num_points, size_x, size_y]
+
+    # all points that will be finally used
+    masked_points = ((extended_x >= 0)
+                     * (extended_x <= image_height - 1)
+                     * (extended_y >= 0)
+                     * (extended_y <= image_width - 1)
+                     * (value >= 0))
+
+    true_extended_x = extended_x
+    true_extended_y = extended_y
+
+    # to prevent error
+    extended_x = (extended_x % image_height)
+    extended_y = (extended_y % image_width)
+
+    # [batch, num_points, size_x, size_y]
+    distance = torch.abs((extended_x - _x.unsqueeze(2).unsqueeze(3))
+                         * (extended_y - _y.unsqueeze(2).unsqueeze(3)))
+    weight = (masked_points.float()
+          * (1 / (value + epsilon)))  # [batch, num_points, size_x, size_y]
+    weighted_value = value * weight
+
+    weight = weight.view([batch, -1])
+    weighted_value = weighted_value.view([batch, -1])
+
+    coordinates = (extended_x.view([batch, -1]) * image_width) + extended_y.view(
+        [batch, -1])
+    coord_max = image_height * image_width
+    true_coordinates = (true_extended_x.view([batch, -1]) * image_width) + true_extended_y.view(
+        [batch, -1])
+    true_coordinates[~masked_points.view([batch, -1])] = coord_max
+    weight_scattered = torch.zeros(
+        [batch, image_width * image_height],
+        device=depth.device).scatter_add(1, coordinates.long(), weight)
+
+    masked_zero_weight_scattered = (weight_scattered == 0.0)
+    weight_scattered += masked_zero_weight_scattered.float()
+
+    weighed_value_scattered = torch.zeros(
+        [batch, image_width * image_height],
+        device=depth.device).scatter_add(1, coordinates.long(), weighted_value)
+
+    return weighed_value_scattered,  weight_scattered
+
+
+def points2depth(points, image_height, image_width, size_x=4, size_y=4):
+    """
+    :param points: [B, num_points, 3]
+    :param image_width:
+    :param image_height:
+    :param size_x:
+    :param size_y:
+    :return:
+        depth_recovered: [B, image_width, image_height]
+    """
+
+    epsilon = torch.tensor([1e-12], requires_grad=False, device=points.device)
+    # epsilon not needed, kept here to ensure exact replication of old version
+    coord_x = (points[:, :, 0] / (points[:, :, 2] + epsilon)) * (image_width / image_height)  # [batch, num_points]
+    coord_y = (points[:, :, 1] / (points[:, :, 2] + epsilon))  # [batch, num_points]
+
+    batch, total_points, _ = points.size()
+    depth = points[:, :, 2]  # [batch, num_points]
+    # pdb.set_trace()
+    _x = ((coord_x + 1) * image_height) / 2
+    _y = ((coord_y + 1) * image_width) / 2
+
+    weighed_value_scattered, weight_scattered = distribute(
+        depth=depth,
+        _x=_x,
+        _y=_y,
+        size_x=size_x,
+        size_y=size_y,
+        image_height=image_height,
+        image_width=image_width)
+
+    depth_recovered = (weighed_value_scattered / weight_scattered).view([
+        batch, image_height, image_width
+    ])
+
+    return depth_recovered
+
+
+# source: https://discuss.pytorch.org/t/batched-index-select/9115/6
+def batched_index_select(inp, dim, index):
+    """
+    input: B x * x ... x *
+    dim: 0 < scalar
+    index: B x M
+    """
+    views = [inp.shape[0]] + \
+        [1 if i != dim else -1 for i in range(1, len(inp.shape))]
+    expanse = list(inp.shape)
+    expanse[0] = -1
+    expanse[dim] = -1
+    index = index.view(views).expand(expanse)
+    return torch.gather(inp, dim, index)
+
+
+def point_fea_img_fea(point_fea, point_coo, h, w):
+    """
+    each point_coo is of the form (x*w + h). points not in the canvas are removed
+    :param point_fea: [batch_size, num_points, feat_size]
+    :param point_coo: [batch_size, num_points]
+    :return:
+    """
+    assert len(point_fea.shape) == 3
+    assert len(point_coo.shape) == 2
+    assert point_fea.shape[0:2] == point_coo.shape
+
+    coo_max = ((h - 1) * w) + (w - 1)
+    mask_point_coo = (point_coo >= 0) * (point_coo <= coo_max)
+    point_coo *= mask_point_coo.float()
+    point_fea *= mask_point_coo.float().unsqueeze(-1)
+
+    bs, _, fs = point_fea.shape
+    point_coo = point_coo.unsqueeze(2).repeat([1, 1, fs])
+    img_fea = torch.zeros([bs, h * w, fs], device=point_fea.device).scatter_add(1, point_coo.long(), point_fea)
+
+    return img_fea
+
+
+def distribute_img_fea_points(img_fea, point_coord):
+    """
+    :param img_fea: [B, C, H, W]
+    :param point_coord: [B, num_points], each coordinate  is a scalar value given by (x * W) + y
+    :return
+        point_fea: [B, num_points, C], for points with coordinates outside the image, we return 0
+    """
+    B, C, H, W = list(img_fea.size())
+    img_fea = img_fea.permute(0, 2, 3, 1).view([B, H*W, C])
+
+    coord_max = ((H - 1) * W) + (W - 1)
+    mask_point_coord = (point_coord >= 0) * (point_coord <= coord_max)
+    mask_point_coord = mask_point_coord.float()
+    point_coord = mask_point_coord * point_coord
+    point_fea = batched_index_select(
+        inp=img_fea,
+        dim=1,
+        index=point_coord.long())
+    point_fea = mask_point_coord.unsqueeze(-1) * point_fea
+    return point_fea
+
+
+class PCViews:
+    """For creating images from PC based on the view information. Faster as the
+    repeated operations are done only once whie initialization.
+    """
+
+    def __init__(self):
+        _views = np.asarray([
+            [[0 * np.pi / 2, 0, np.pi / 2], [0, 0, TRANS]],
+            [[1 * np.pi / 2, 0, np.pi / 2], [0, 0, TRANS]],
+            [[2 * np.pi / 2, 0, np.pi / 2], [0, 0, TRANS]],
+            [[3 * np.pi / 2, 0, np.pi / 2], [0, 0, TRANS]],
+            [[0, -np.pi / 2, np.pi / 2], [0, 0, TRANS]],
+            [[0, np.pi / 2, np.pi / 2], [0, 0, TRANS]]])
+        self.num_views = 6
+        angle = torch.tensor(_views[:, 0, :]).float().cuda()
+        self.rot_mat = euler2mat(angle).transpose(1, 2)
+        self.translation = torch.tensor(_views[:, 1, :]).float().cuda()
+        self.translation = self.translation.unsqueeze(1)
+
+    def get_img(self, points):
+        """Get image based on the prespecified specifications.
+
+        Args:
+            points (torch.tensor): of size [B, _, 3]
+        Returns:
+            img (torch.tensor): of size [B * self.num_views, RESOLUTION,
+                RESOLUTION]
+        """
+        b, _, _ = points.shape
+        v = self.translation.shape[0]
+
+        _points = self.point_transform(
+            points=torch.repeat_interleave(points, v, dim=0),
+            rot_mat=self.rot_mat.repeat(b, 1, 1),
+            translation=self.translation.repeat(b, 1, 1))
+
+        img = points2depth(
+            points=_points,
+            image_height=RESOLUTION,
+            image_width=RESOLUTION,
+            size_x=1,
+            size_y=1,
+        )
+        return img
+
+    @staticmethod
+    def point_transform(points, rot_mat, translation):
+        """
+        :param points: [batch, num_points, 3]
+        :param rot_mat: [batch, 3]
+        :param translation: [batch, 1, 3]
+        :return:
+        """
+        rot_mat = rot_mat.to(points.device)
+        translation = translation.to(points.device)
+        points = torch.matmul(points, rot_mat)
+        points = points - translation
+        return points

openpoints/models/build.py

@@ -0,0 +1,13 @@
+from openpoints.utils import registry
+MODELS = registry.Registry('models')
+
+
+def build_model_from_cfg(cfg, **kwargs):
+    """
+    Build a model, defined by `NAME`.
+    Args:
+        cfg (eDICT): 
+    Returns:
+        Model: a constructed model specified by NAME.
+    """
+    return MODELS.build(cfg, **kwargs)

openpoints/models/classification/__init__.py

@@ -0,0 +1,5 @@
+"""
+Author: PointNeXt
+
+"""
+from .cls_base import BaseCls, ClsHead

openpoints/models/classification/cls_base.py

@@ -0,0 +1,136 @@
+import torch
+import torch.nn as nn
+import logging
+from typing import List
+from ..layers import create_linearblock
+from ...utils import get_missing_parameters_message, get_unexpected_parameters_message
+from ..build import MODELS, build_model_from_cfg
+from ...loss import build_criterion_from_cfg
+from ...utils import load_checkpoint
+
+
+@MODELS.register_module()
+class BaseCls(nn.Module):
+    def __init__(self,
+                 encoder_args=None,
+                 cls_args=None,
+                 criterion_args=None,
+                 **kwargs):
+        super().__init__()
+        self.encoder = build_model_from_cfg(encoder_args)
+
+        if cls_args is not None:
+            in_channels = self.encoder.out_channels if hasattr(self.encoder, 'out_channels') else cls_args.get('in_channels', None)
+            cls_args.in_channels = in_channels
+            self.prediction = build_model_from_cfg(cls_args)
+        else:
+            self.prediction = nn.Identity()
+        self.criterion = build_criterion_from_cfg(criterion_args) if criterion_args is not None else None
+
+    def forward(self, data):
+        global_feat = self.encoder.forward_cls_feat(data)
+        return self.prediction(global_feat)
+
+    def get_loss(self, pred, gt, inputs=None):
+        return self.criterion(pred, gt.long())
+
+    def get_logits_loss(self, data, gt):
+        logits = self.forward(data)
+        return logits, self.criterion(logits, gt.long())
+
+
+@MODELS.register_module()
+class DistillCls(BaseCls):
+    def __init__(self,
+                 encoder_args=None,
+                 cls_args=None,
+                 distill_args=None,
+                 criterion_args=None,
+                 **kwargs):
+        super().__init__(encoder_args, cls_args, criterion_args)
+        self.distill = encoder_args.get('distill', True)
+        in_channels = self.encoder.distill_channels
+        distill_args.distill_head_args.in_channels = in_channels
+        self.dist_head = build_model_from_cfg(distill_args.distill_head_args)
+        self.dist_model = build_model_from_cfg(distill_args).cuda()
+        load_checkpoint(self.dist_model, distill_args.pretrained_path)
+        self.dist_model.eval()
+
+    def forward(self, p0, f0=None):
+        if hasattr(p0, 'keys'):
+            p0, f0 = p0['pos'], p0['x']
+        if self.distill and self.training:
+            global_feat, distill_feature = self.encoder.forward_cls_feat(p0, f0)
+            return self.prediction(global_feat), self.dist_head(distill_feature)
+        else:
+            global_feat = self.encoder.forward_cls_feat(p0, f0)
+            return self.prediction(global_feat)
+
+    def get_loss(self, pred, gt, inputs):
+        return self.criterion(inputs, pred, gt.long(), self.dist_model)
+
+    def get_logits_loss(self, data, gt):
+        logits, dist_logits = self.forward(data)
+        return logits, self.criterion(data, [logits, dist_logits], gt.long(), self.dist_model)
+
+
+@MODELS.register_module()
+class ClsHead(nn.Module):
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int,
+                 mlps: List[int]=[256],
+                 norm_args: dict=None,
+                 act_args: dict={'act': 'relu'},
+                 dropout: float=0.5,
+                 global_feat: str=None,
+                 point_dim: int=2,
+                 **kwargs
+                 ):
+        """A general classification head. supports global pooling and [CLS] token
+        Args:
+            num_classes (int): class num
+            in_channels (int): input channels size
+            mlps (List[int], optional): channel sizes for hidden layers. Defaults to [256].
+            norm_args (dict, optional): dict of configuration for normalization. Defaults to None.
+            act_args (_type_, optional): dict of configuration for activation. Defaults to {'act': 'relu'}.
+            dropout (float, optional): use dropout when larger than 0. Defaults to 0.5.
+            cls_feat (str, optional): preprocessing input features to obtain global feature.
+                                      $\eg$ cls_feat='max,avg' means use the concatenateion of maxpooled and avgpooled features.
+                                      Defaults to None, which means the input feautre is the global feature
+        Returns:
+            logits: (B, num_classes, N)
+        """
+        super().__init__()
+        if kwargs:
+            logging.warning(f"kwargs: {kwargs} are not used in {__class__.__name__}")
+        self.global_feat = global_feat.split(',') if global_feat is not None else None
+        self.point_dim = point_dim
+        in_channels = len(self.global_feat) * in_channels if global_feat is not None else in_channels
+        if mlps is not None:
+            mlps = [in_channels] + mlps + [num_classes]
+        else:
+            mlps = [in_channels, num_classes]
+
+        heads = []
+        for i in range(len(mlps) - 2):
+            heads.append(create_linearblock(mlps[i], mlps[i + 1],
+                                            norm_args=norm_args,
+                                            act_args=act_args))
+            if dropout:
+                heads.append(nn.Dropout(dropout))
+        heads.append(create_linearblock(mlps[-2], mlps[-1], act_args=None))
+        self.head = nn.Sequential(*heads)
+
+
+    def forward(self, end_points):
+        if self.global_feat is not None:
+            global_feats = []
+            for preprocess in self.global_feat:
+                if 'max' in preprocess:
+                    global_feats.append(torch.max(end_points, dim=self.point_dim, keepdim=False)[0])
+                elif preprocess in ['avg', 'mean']:
+                    global_feats.append(torch.mean(end_points, dim=self.point_dim, keepdim=False))
+            end_points = torch.cat(global_feats, dim=1)
+        logits = self.head(end_points)
+        return logits