indoor_loading.py 4.13 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
3
4
5
6
7
8
import os.path as osp

import numpy as np

from mmdet.datasets.registry import PIPELINES


@PIPELINES.register_module
liyinhao's avatar
liyinhao committed
9
10
class LoadPointsFromFile(object):
    """Load Points From File.
liyinhao's avatar
liyinhao committed
11

liyinhao's avatar
liyinhao committed
12
    Load sunrgbd and scannet points from file.
liyinhao's avatar
liyinhao committed
13
14
15
16
17
18
19
20

    Args:
        name (str): scannet or sunrgbd.
        use_color (bool): Whether to use color.
        use_height (bool): Whether to use height.
        mean_color (List[float]): Mean color of the point cloud.
    """

liyinhao's avatar
liyinhao committed
21
    def __init__(self, use_color, use_height, mean_color):
liyinhao's avatar
liyinhao committed
22
23
24
25
26
27
28
        self.use_color = use_color
        self.use_height = use_height
        self.mean_color = mean_color

    def __call__(self, results):
        data_path = results.get('data_path', None)
        info = results.get('info', None)
liyinhao's avatar
liyinhao committed
29
30
        name = 'scannet' if info.get('image', None) is None else 'sunrgbd'
        if name == 'scannet':
liyinhao's avatar
liyinhao committed
31
32
            scan_name = info['point_cloud']['lidar_idx']
            point_cloud = self._get_lidar(scan_name, data_path)
liyinhao's avatar
liyinhao committed
33
        else:
liyinhao's avatar
liyinhao committed
34
            point_cloud = np.load(
liyinhao's avatar
liyinhao committed
35
                osp.join(data_path, 'lidar',
liyinhao's avatar
liyinhao committed
36
                         '%06d.npz' % info['point_cloud']['lidar_idx']))['pc']
liyinhao's avatar
liyinhao committed
37

liyinhao's avatar
liyinhao committed
38
        if not self.use_color:
liyinhao's avatar
liyinhao committed
39
            if name == 'scannet':
liyinhao's avatar
liyinhao committed
40
                pcl_color = point_cloud[:, 3:6]
liyinhao's avatar
liyinhao committed
41
            point_cloud = point_cloud[:, 0:3]
liyinhao's avatar
liyinhao committed
42
        else:
liyinhao's avatar
liyinhao committed
43
            if name == 'scannet':
liyinhao's avatar
liyinhao committed
44
                pcl_color = point_cloud[:, 3:6]
liyinhao's avatar
liyinhao committed
45
46
47
            point_cloud = point_cloud[:, 0:6]
            point_cloud[:, 3:] = (point_cloud[:, 3:] -
                                  np.array(self.mean_color)) / 256.0
liyinhao's avatar
liyinhao committed
48

liyinhao's avatar
liyinhao committed
49
        if self.use_height:
liyinhao's avatar
liyinhao committed
50
51
52
53
54
            floor_height = np.percentile(point_cloud[:, 2], 0.99)
            height = point_cloud[:, 2] - floor_height
            point_cloud = np.concatenate(
                [point_cloud, np.expand_dims(height, 1)], 1)
        results['point_cloud'] = point_cloud
liyinhao's avatar
liyinhao committed
55
        if name == 'scannet':
liyinhao's avatar
liyinhao committed
56
57
58
59
60
            results['pcl_color'] = pcl_color
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
liyinhao's avatar
liyinhao committed
61
62
63
        repr_str += '(use_height={})'.format(self.use_height)
        repr_str += '(use_color={}'.format(self.use_color)
        repr_str += '(mean_color={})'.format(self.mean_color)
liyinhao's avatar
liyinhao committed
64
65
66
67
68
69
70
        return repr_str

    def _get_lidar(self, scan_name, data_path):
        lidar_file = osp.join(data_path, scan_name + '_vert.npy')
        assert osp.exists(lidar_file)
        return np.load(lidar_file)

liyinhao's avatar
liyinhao committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

@PIPELINES.register_module
class LoadAnnotations3D(object):
    """Load Annotations3D.

    Load sunrgbd and scannet annotations.

    Args:
        name (str): scannet or sunrgbd.
    """

    def __init__(self):
        pass

    def __call__(self, results):
        data_path = results.get('data_path', None)
        info = results.get('info', None)
        if info['annos']['gt_num'] != 0:
            gt_boxes = info['annos']['gt_boxes_upright_depth']
            gt_classes = info['annos']['class'].reshape(-1, 1)
            gt_boxes_mask = np.ones_like(gt_classes)
        else:
            gt_boxes = np.zeros((1, 6), dtype=np.float32)
            gt_classes = np.zeros((1, 1))
            gt_boxes_mask = np.zeros((1, 1))
        name = 'scannet' if info.get('image', None) is None else 'sunrgbd'

        if name == 'scannet':
            scan_name = info['point_cloud']['lidar_idx']
            instance_labels = self._get_instance_label(scan_name, data_path)
            semantic_labels = self._get_semantic_label(scan_name, data_path)
            results['instance_labels'] = instance_labels
            results['semantic_labels'] = semantic_labels

        results['gt_boxes'] = gt_boxes
        results['gt_classes'] = gt_classes
        results['gt_boxes_mask'] = gt_boxes_mask
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        return repr_str

liyinhao's avatar
liyinhao committed
114
115
116
117
118
119
120
121
122
    def _get_instance_label(self, scan_name, data_path):
        ins_file = osp.join(data_path, scan_name + '_ins_label.npy')
        assert osp.exists(ins_file)
        return np.load(ins_file)

    def _get_semantic_label(self, scan_name, data_path):
        sem_file = osp.join(data_path, scan_name + '_sem_label.npy')
        assert osp.exists(sem_file)
        return np.load(sem_file)