sunrgbd_dataset.py 1.57 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
3
4
5
import os.path as osp

import numpy as np

from mmdet.datasets import DATASETS
6
from .indoor_base_dataset import IndoorBaseDataset
liyinhao's avatar
liyinhao committed
7
8
9


@DATASETS.register_module()
10
class SUNRGBDDataset(IndoorBaseDataset):
liyinhao's avatar
liyinhao committed
11

liyinhao's avatar
liyinhao committed
12
13
14
15
16
17
18
    CLASSES = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
               'night_stand', 'bookshelf', 'bathtub')

    def __init__(self,
                 root_path,
                 ann_file,
                 pipeline=None,
liyinhao's avatar
liyinhao committed
19
                 classes=None,
liyinhao's avatar
liyinhao committed
20
21
                 test_mode=False,
                 with_label=True):
liyinhao's avatar
liyinhao committed
22
23
        super().__init__(root_path, ann_file, pipeline, classes, test_mode,
                         with_label)
liyinhao's avatar
liyinhao committed
24
25

    def _get_pts_filename(self, sample_idx):
26
        pts_filename = osp.join(self.root_path, 'lidar',
27
                                f'{sample_idx:06d}.npy')
liyinhao's avatar
liyinhao committed
28
29
30
31
        return pts_filename

    def _get_ann_info(self, index, sample_idx):
        # Use index to get the annos, thus the evalhook could also use this api
liyinhao's avatar
liyinhao committed
32
        info = self.data_infos[index]
liyinhao's avatar
liyinhao committed
33
34
35
        if info['annos']['gt_num'] != 0:
            gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']  # k, 6
            gt_labels = info['annos']['class']
36
            gt_bboxes_3d_mask = np.ones_like(gt_labels, dtype=np.bool)
liyinhao's avatar
liyinhao committed
37
38
        else:
            gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
39
40
            gt_labels = np.zeros(1, dtype=np.bool)
            gt_bboxes_3d_mask = np.zeros(1, dtype=np.bool)
liyinhao's avatar
liyinhao committed
41
42
43
44
45
46

        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
            gt_labels=gt_labels,
            gt_bboxes_3d_mask=gt_bboxes_3d_mask)
        return anns_results