encoder_decoder.py 21.7 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
from typing import List, Tuple
3

4
5
import numpy as np
import torch
6
from torch import Tensor
7
8
9
from torch import nn as nn
from torch.nn import functional as F

10
from mmdet3d.registry import MODELS
zhangshilong's avatar
zhangshilong committed
11
12
13
from mmdet3d.utils import ConfigType, OptConfigType, OptMultiConfig
from ...structures.det3d_data_sample import OptSampleList, SampleList
from ..utils import add_prefix
14
15
16
from .base import Base3DSegmentor


17
@MODELS.register_module()
18
19
20
21
22
class EncoderDecoder3D(Base3DSegmentor):
    """3D Encoder Decoder segmentors.

    EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
    Note that auxiliary_head is only used for deep supervision during training,
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
    which could be dumped during inference.

    1. The ``loss`` method is used to calculate the loss of model,
    which includes two steps: (1) Extracts features to obtain the feature maps
    (2) Call the decode head loss function to forward decode head model and
    calculate losses.

    .. code:: text

     loss(): extract_feat() -> _decode_head_forward_train() -> _auxiliary_head_forward_train (optional)
     _decode_head_forward_train(): decode_head.loss()
     _auxiliary_head_forward_train(): auxiliary_head.loss (optional)

    2. The ``predict`` method is used to predict segmentation results,
    which includes two steps: (1) Run inference function to obtain the list of
    seg_logits (2) Call post-processing function to obtain list of
    ``SegDataSampel`` including ``pred_sem_seg`` and ``seg_logits``.

    .. code:: text

43
44
45
46
    predict(): inference() -> postprocess_result()
    inference(): whole_inference()/slide_inference()
    whole_inference()/slide_inference(): encoder_decoder()
    encoder_decoder(): extract_feat() -> decode_head.predict()
47
48
49
50
51
52
53

    4 The ``_forward`` method is used to output the tensor by running the model,
    which includes two steps: (1) Extracts features to obtain the feature maps
    (2)Call the decode head forward function to forward decode head model.

    .. code:: text

54
    _forward(): extract_feat() -> _decode_head.forward()
55
56
57
58
59
60
61
62
63
64
65
66
67

    Args:

        backbone (ConfigType): The config for the backnone of segmentor.
        decode_head (ConfigType): The config for the decode head of segmentor.
        neck (OptConfigType): The config for the neck of segmentor.
            Defaults to None.
        auxiliary_head (OptConfigType): The config for the auxiliary head of
            segmentor. Defaults to None.
        loss_regularization (OptiConfigType): The config for the regularization
            loass. Defaults to None.
        train_cfg (OptConfigType): The config for training. Defaults to None.
        test_cfg (OptConfigType): The config for testing. Defaults to None.
68
69
70
71
        data_preprocessor (OptConfigType): The pre-process config of
            :class:`BaseDataPreprocessor`. Defaults to None.
        init_cfg (OptMultiConfig): The weight initialized config for
            :class:`BaseModule`. Defaults to None.
72
    """  # noqa: E501
73
74

    def __init__(self,
75
76
77
78
79
80
81
82
                 backbone: ConfigType,
                 decode_head: ConfigType,
                 neck: OptConfigType = None,
                 auxiliary_head: OptConfigType = None,
                 loss_regularization: OptConfigType = None,
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None,
                 data_preprocessor: OptConfigType = None,
83
                 init_cfg: OptMultiConfig = None) -> None:
84
85
        super(EncoderDecoder3D, self).__init__(
            data_preprocessor=data_preprocessor, init_cfg=init_cfg)
86
        self.backbone = MODELS.build(backbone)
87
        if neck is not None:
88
            self.neck = MODELS.build(neck)
89
90
        self._init_decode_head(decode_head)
        self._init_auxiliary_head(auxiliary_head)
91
        self._init_loss_regularization(loss_regularization)
92
93
94

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
95

96
97
98
        assert self.with_decode_head, \
            '3D EncoderDecoder Segmentor should have a decode_head'

99
    def _init_decode_head(self, decode_head: ConfigType) -> None:
100
        """Initialize ``decode_head``"""
101
        self.decode_head = MODELS.build(decode_head)
102
103
        self.num_classes = self.decode_head.num_classes

104
    def _init_auxiliary_head(self, auxiliary_head: ConfigType) -> None:
105
106
107
108
109
        """Initialize ``auxiliary_head``"""
        if auxiliary_head is not None:
            if isinstance(auxiliary_head, list):
                self.auxiliary_head = nn.ModuleList()
                for head_cfg in auxiliary_head:
110
                    self.auxiliary_head.append(MODELS.build(head_cfg))
111
            else:
112
                self.auxiliary_head = MODELS.build(auxiliary_head)
113

114
115
    def _init_loss_regularization(self,
                                  loss_regularization: ConfigType) -> None:
116
117
118
119
120
        """Initialize ``loss_regularization``"""
        if loss_regularization is not None:
            if isinstance(loss_regularization, list):
                self.loss_regularization = nn.ModuleList()
                for loss_cfg in loss_regularization:
121
                    self.loss_regularization.append(MODELS.build(loss_cfg))
122
            else:
123
                self.loss_regularization = MODELS.build(loss_regularization)
124

125
    def extract_feat(self, batch_inputs: Tensor) -> Tensor:
126
        """Extract features from points."""
127
        x = self.backbone(batch_inputs)
128
129
130
131
        if self.with_neck:
            x = self.neck(x)
        return x

132
133
    def encode_decode(self, batch_inputs: Tensor,
                      batch_input_metas: List[dict]) -> Tensor:
134
135
136
137
        """Encode points with backbone and decode into a semantic segmentation
        map of the same size as input.

        Args:
138
            batch_input (torch.Tensor): Input point cloud sample
139
            batch_input_metas (list[dict]): Meta information of each sample.
140
141
142
143

        Returns:
            torch.Tensor: Segmentation logits of shape [B, num_classes, N].
        """
144
        x = self.extract_feat(batch_inputs)
145
146
147
        seg_logits = self.decode_head.predict(x, batch_input_metas,
                                              self.test_cfg)
        return seg_logits
148

149
150
    def _decode_head_forward_train(self, batch_inputs_dict: dict,
                                   batch_data_samples: SampleList) -> dict:
151
152
153
        """Run forward function and calculate loss for decode head in
        training."""
        losses = dict()
154
155
        loss_decode = self.decode_head.loss(batch_inputs_dict,
                                            batch_data_samples, self.train_cfg)
156
157
158
159

        losses.update(add_prefix(loss_decode, 'decode'))
        return losses

160
161
162
163
164
    def _auxiliary_head_forward_train(
        self,
        batch_inputs_dict: dict,
        batch_data_samples: SampleList,
    ) -> dict:
165
166
167
168
169
        """Run forward function and calculate loss for auxiliary head in
        training."""
        losses = dict()
        if isinstance(self.auxiliary_head, nn.ModuleList):
            for idx, aux_head in enumerate(self.auxiliary_head):
170
171
                loss_aux = aux_head.loss(batch_inputs_dict, batch_data_samples,
                                         self.train_cfg)
172
173
                losses.update(add_prefix(loss_aux, f'aux_{idx}'))
        else:
174
175
176
            loss_aux = self.auxiliary_head.loss(batch_inputs_dict,
                                                batch_data_samples,
                                                self.train_cfg)
177
178
179
180
            losses.update(add_prefix(loss_aux, 'aux'))

        return losses

181
    def _loss_regularization_forward_train(self) -> dict:
182
183
184
185
186
187
188
189
190
191
192
193
194
195
        """Calculate regularization loss for model weight in training."""
        losses = dict()
        if isinstance(self.loss_regularization, nn.ModuleList):
            for idx, regularize_loss in enumerate(self.loss_regularization):
                loss_regularize = dict(
                    loss_regularize=regularize_loss(self.modules()))
                losses.update(add_prefix(loss_regularize, f'regularize_{idx}'))
        else:
            loss_regularize = dict(
                loss_regularize=self.loss_regularization(self.modules()))
            losses.update(add_prefix(loss_regularize, 'regularize'))

        return losses

196
197
198
    def loss(self, batch_inputs_dict: dict,
             batch_data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.
199
200

        Args:
201
202
203
204
205
            batch_inputs_dict (dict): Input sample dict which
                includes 'points' and 'imgs' keys.

                - points (list[torch.Tensor]): Point cloud of each sample.
                - imgs (torch.Tensor, optional): Image tensor has shape
206
                  (B, C, H, W).
207
208
209
            batch_data_samples (list[:obj:`Det3DDataSample`]): The det3d
                data samples. It usually includes information such
                as `metainfo` and `gt_pts_sem_seg`.
210
211

        Returns:
212
            dict[str, Tensor]: a dictionary of loss components.
213
214
215
        """

        # extract features using backbone
216
217
        points = torch.stack(batch_inputs_dict['points'])
        x = self.extract_feat(points)
218
219
220

        losses = dict()

221
        loss_decode = self._decode_head_forward_train(x, batch_data_samples)
222
223
224
225
        losses.update(loss_decode)

        if self.with_auxiliary_head:
            loss_aux = self._auxiliary_head_forward_train(
226
                x, batch_data_samples)
227
228
            losses.update(loss_aux)

229
230
231
232
        if self.with_regularization_loss:
            loss_regularize = self._loss_regularization_forward_train()
            losses.update(loss_regularize)

233
234
235
236
        return losses

    @staticmethod
    def _input_generation(coords,
237
238
239
                          patch_center: Tensor,
                          coord_max: Tensor,
                          feats: Tensor,
240
                          use_normalized_coord: bool = False) -> Tensor:
241
242
        """Generating model input.

243
        Generate input by subtracting patch center and adding additional
244
245
246
247
248
249
250
            features. Currently support colors and normalized xyz as features.

        Args:
            coords (torch.Tensor): Sampled 3D point coordinate of shape [S, 3].
            patch_center (torch.Tensor): Center coordinate of the patch.
            coord_max (torch.Tensor): Max coordinate of all 3D points.
            feats (torch.Tensor): Features of sampled points of shape [S, C].
251
            use_normalized_coord (bool, optional): Whether to use normalized
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
                xyz as additional features. Defaults to False.

        Returns:
            torch.Tensor: The generated input data of shape [S, 3+C'].
        """
        # subtract patch center, the z dimension is not centered
        centered_coords = coords.clone()
        centered_coords[:, 0] -= patch_center[0]
        centered_coords[:, 1] -= patch_center[1]

        # normalized coordinates as extra features
        if use_normalized_coord:
            normalized_coord = coords / coord_max
            feats = torch.cat([feats, normalized_coord], dim=1)

        points = torch.cat([centered_coords, feats], dim=1)

        return points

    def _sliding_patch_generation(self,
272
273
274
275
276
                                  points: Tensor,
                                  num_points: int,
                                  block_size: float,
                                  sample_rate: float = 0.5,
                                  use_normalized_coord: bool = False,
277
                                  eps: float = 1e-3) -> Tuple[Tensor, Tensor]:
278
279
280
281
282
283
284
285
286
287
288
        """Sampling points in a sliding window fashion.

        First sample patches to cover all the input points.
        Then sample points in each patch to batch points of a certain number.

        Args:
            points (torch.Tensor): Input points of shape [N, 3+C].
            num_points (int): Number of points to be sampled in each patch.
            block_size (float, optional): Size of a patch to sample.
            sample_rate (float, optional): Stride used in sliding patch.
                Defaults to 0.5.
289
            use_normalized_coord (bool, optional): Whether to use normalized
290
291
                xyz as additional features. Defaults to False.
            eps (float, optional): A value added to patch boundary to guarantee
292
                points coverage. Defaults to 1e-3.
293
294

        Returns:
295
            tuple:
296

297
                - patch_points (torch.Tensor): Points of different patches of
298
                  shape [K, N, 3+C].
299
                - patch_idxs (torch.Tensor): Index of each point in
300
                  `patch_points`, of shape [K, N].
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
        """
        device = points.device
        # we assume the first three dims are points' 3D coordinates
        # and the rest dims are their per-point features
        coords = points[:, :3]
        feats = points[:, 3:]

        coord_max = coords.max(0)[0]
        coord_min = coords.min(0)[0]
        stride = block_size * sample_rate
        num_grid_x = int(
            torch.ceil((coord_max[0] - coord_min[0] - block_size) /
                       stride).item() + 1)
        num_grid_y = int(
            torch.ceil((coord_max[1] - coord_min[1] - block_size) /
                       stride).item() + 1)

        patch_points, patch_idxs = [], []
        for idx_y in range(num_grid_y):
            s_y = coord_min[1] + idx_y * stride
            e_y = torch.min(s_y + block_size, coord_max[1])
            s_y = e_y - block_size
            for idx_x in range(num_grid_x):
                s_x = coord_min[0] + idx_x * stride
                e_x = torch.min(s_x + block_size, coord_max[0])
                s_x = e_x - block_size

                # extract points within this patch
                cur_min = torch.tensor([s_x, s_y, coord_min[2]]).to(device)
                cur_max = torch.tensor([e_x, e_y, coord_max[2]]).to(device)
                cur_choice = ((coords >= cur_min - eps) &
                              (coords <= cur_max + eps)).all(dim=1)

                if not cur_choice.any():  # no points in this patch
                    continue

                # sample points in this patch to multiple batches
                cur_center = cur_min + block_size / 2.0
                point_idxs = torch.nonzero(cur_choice, as_tuple=True)[0]
                num_batch = int(np.ceil(point_idxs.shape[0] / num_points))
                point_size = int(num_batch * num_points)
                replace = point_size > 2 * point_idxs.shape[0]
                num_repeat = point_size - point_idxs.shape[0]
                if replace:  # duplicate
                    point_idxs_repeat = point_idxs[torch.randint(
                        0, point_idxs.shape[0],
                        size=(num_repeat, )).to(device)]
                else:
                    point_idxs_repeat = point_idxs[torch.randperm(
                        point_idxs.shape[0])[:num_repeat]]

                choices = torch.cat([point_idxs, point_idxs_repeat], dim=0)
                choices = choices[torch.randperm(choices.shape[0])]

                # construct model input
                point_batches = self._input_generation(
                    coords[choices],
                    cur_center,
                    coord_max,
                    feats[choices],
                    use_normalized_coord=use_normalized_coord)

                patch_points.append(point_batches)
                patch_idxs.append(choices)

        patch_points = torch.cat(patch_points, dim=0)
        patch_idxs = torch.cat(patch_idxs, dim=0)

        # make sure all points are sampled at least once
        assert torch.unique(patch_idxs).shape[0] == points.shape[0], \
            'some points are not sampled in sliding inference'

        return patch_points, patch_idxs

375
    def slide_inference(self, point: Tensor, img_meta: List[dict],
376
                        rescale: bool) -> Tensor:
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
        """Inference by sliding-window with overlap.

        Args:
            point (torch.Tensor): Input points of shape [N, 3+C].
            img_meta (dict): Meta information of input sample.
            rescale (bool): Whether transform to original number of points.
                Will be used for voxelization based segmentors.

        Returns:
            Tensor: The output segmentation map of shape [num_classes, N].
        """
        num_points = self.test_cfg.num_points
        block_size = self.test_cfg.block_size
        sample_rate = self.test_cfg.sample_rate
        use_normalized_coord = self.test_cfg.use_normalized_coord
        batch_size = self.test_cfg.batch_size * num_points

        # patch_points is of shape [K*N, 3+C], patch_idxs is of shape [K*N]
        patch_points, patch_idxs = self._sliding_patch_generation(
            point, num_points, block_size, sample_rate, use_normalized_coord)
        feats_dim = patch_points.shape[1]
        seg_logits = []  # save patch predictions

        for batch_idx in range(0, patch_points.shape[0], batch_size):
            batch_points = patch_points[batch_idx:batch_idx + batch_size]
            batch_points = batch_points.view(-1, num_points, feats_dim)
            # batch_seg_logit is of shape [B, num_classes, N]
            batch_seg_logit = self.encode_decode(batch_points, img_meta)
            batch_seg_logit = batch_seg_logit.transpose(1, 2).contiguous()
            seg_logits.append(batch_seg_logit.view(-1, self.num_classes))

        # aggregate per-point logits by indexing sum and dividing count
        seg_logits = torch.cat(seg_logits, dim=0)  # [K*N, num_classes]
        expand_patch_idxs = patch_idxs.unsqueeze(1).repeat(1, self.num_classes)
        preds = point.new_zeros((point.shape[0], self.num_classes)).\
            scatter_add_(dim=0, index=expand_patch_idxs, src=seg_logits)
        count_mat = torch.bincount(patch_idxs)
        preds = preds / count_mat[:, None]

        # TODO: if rescale and voxelization segmentor

        return preds.transpose(0, 1)  # to [num_classes, K*N]

420
    def whole_inference(self, points: Tensor, input_metas: List[dict],
421
                        rescale: bool) -> Tensor:
422
        """Inference with full scene (one forward pass without sliding)."""
423
        seg_logit = self.encode_decode(points, input_metas)
424
425
426
        # TODO: if rescale and voxelization segmentor
        return seg_logit

427
    def inference(self, points: Tensor, input_metas: List[dict],
428
                  rescale: bool) -> Tensor:
429
430
431
432
        """Inference with slide/whole style.

        Args:
            points (torch.Tensor): Input points of shape [B, N, 3+C].
433
            input_metas (list[dict]): Meta information of each sample.
434
435
436
437
438
439
440
441
442
443
            rescale (bool): Whether transform to original number of points.
                Will be used for voxelization based segmentors.

        Returns:
            Tensor: The output segmentation map.
        """
        assert self.test_cfg.mode in ['slide', 'whole']
        if self.test_cfg.mode == 'slide':
            seg_logit = torch.stack([
                self.slide_inference(point, img_meta, rescale)
444
                for point, img_meta in zip(points, input_metas)
445
446
            ], 0)
        else:
447
            seg_logit = self.whole_inference(points, input_metas, rescale)
448
449
450
        output = F.softmax(seg_logit, dim=1)
        return output

451
452
453
454
    def predict(self,
                batch_inputs_dict: dict,
                batch_data_samples: SampleList,
                rescale: bool = True) -> SampleList:
455
456
457
        """Simple test with single scene.

        Args:
458
459
460
461
462
463
464
465
466
            batch_inputs_dict (dict): Input sample dict which
                includes 'points' and 'imgs' keys.

                - points (list[torch.Tensor]): Point cloud of each sample.
                - imgs (torch.Tensor, optional): Image tensor has shape
                    (B, C, H, W).
            batch_data_samples (list[:obj:`Det3DDataSample`]): The det3d
                data samples. It usually includes information such
                as `metainfo` and `gt_pts_sem_seg`.
467
468
469
470
471
472
473
474
475
476
477
478
            rescale (bool): Whether transform to original number of points.
                Will be used for voxelization based segmentors.
                Defaults to True.

        Returns:
            list[dict]: The output prediction result with following keys:

                - semantic_mask (Tensor): Segmentation mask of shape [N].
        """
        # 3D segmentation requires per-point prediction, so it's impossible
        # to use down-sampling to get a batch of scenes with same num_points
        # therefore, we only support testing one scene every time
479
        seg_pred_list = []
480
481
482
483
484
485
486
487
        batch_input_metas = []
        for data_sample in batch_data_samples:
            batch_input_metas.append(data_sample.metainfo)

        points = batch_inputs_dict['points']
        for point, input_meta in zip(points, batch_input_metas):
            seg_prob = self.inference(
                point.unsqueeze(0), [input_meta], rescale)[0]
488
489
490
            seg_map = seg_prob.argmax(0)  # [N]
            # to cpu tensor for consistency with det3d
            seg_map = seg_map.cpu()
491
492
            seg_pred_list.append(seg_map)

493
        return self.postprocess_result(seg_pred_list, batch_data_samples)
494

495
496
497
498
    def _forward(self,
                 batch_inputs_dict: dict,
                 batch_data_samples: OptSampleList = None) -> Tensor:
        """Network forward process.
499
500

        Args:
501
502
            batch_inputs_dict (dict): Input sample dict which
                includes 'points' and 'imgs' keys.
503

504
505
                - points (list[torch.Tensor]): Point cloud of each sample.
                - imgs (torch.Tensor, optional): Image tensor has shape
506
                  (B, C, H, W).
507
508
509
            batch_data_samples (List[:obj:`Det3DDataSample`]): The seg
                data samples. It usually includes information such
                as `metainfo` and `gt_pts_sem_seg`.
510

511
512
        Returns:
            Tensor: Forward output of model without any post-processes.
513
        """
514
515
        points = torch.stack(batch_inputs_dict['points'])
        x = self.extract_feat(points)
516
        return self.decode_head.forward(x)
517
518
519
520

    def aug_test(self, batch_inputs, batch_img_metas):
        """Placeholder for augmentation test."""
        pass