test_indoor_loading.py 3.7 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
import os.path as osp

liyinhao's avatar
liyinhao committed
3
4
import mmcv

liyinhao's avatar
liyinhao committed
5
6
from mmdet3d.datasets.pipelines.indoor_loading import (LoadAnnotations3D,
                                                       LoadPointsFromFile)
liyinhao's avatar
liyinhao committed
7
8


liyinhao's avatar
liyinhao committed
9
def test_load_points_from_file():
10
    sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
liyinhao's avatar
liyinhao committed
11
12
    sunrgbd_load_points_from_file = LoadPointsFromFile(True, [0.5, 0.5, 0.5],
                                                       6)
liyinhao's avatar
liyinhao committed
13
    sunrgbd_results = dict()
liyinhao's avatar
liyinhao committed
14
15
16
17
18
19
20
    data_path = './tests/data/sunrgbd/sunrgbd_trainval'
    sunrgbd_results['data_path'] = data_path
    sunrgbd_info = sunrgbd_info[0]
    scan_name = sunrgbd_info['point_cloud']['lidar_idx']
    sunrgbd_results['info'] = sunrgbd_info
    sunrgbd_results['pts_filename'] = osp.join(data_path, 'lidar',
                                               '%06d.npz' % scan_name)
liyinhao's avatar
liyinhao committed
21
    sunrgbd_results = sunrgbd_load_points_from_file(sunrgbd_results)
liyinhao's avatar
liyinhao committed
22
    sunrgbd_point_cloud = sunrgbd_results.get('points', None)
liyinhao's avatar
liyinhao committed
23
24
25
    assert sunrgbd_point_cloud.shape == (1000, 4)

    scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
liyinhao's avatar
liyinhao committed
26
    scannet_load_data = LoadPointsFromFile(True, [0.5, 0.5, 0.5])
liyinhao's avatar
liyinhao committed
27
    scannet_results = dict()
liyinhao's avatar
liyinhao committed
28
29
30
31
32
33
34
    data_path = './tests/data/scannet/scannet_train_instance_data'
    scannet_results['data_path'] = data_path
    scannet_info = scannet_info[0]
    scan_name = scannet_info['point_cloud']['lidar_idx']
    scannet_results['info'] = scannet_info
    scannet_results['pts_filename'] = osp.join(data_path,
                                               scan_name + '_vert.npy')
liyinhao's avatar
liyinhao committed
35
    scannet_results = scannet_load_data(scannet_results)
liyinhao's avatar
liyinhao committed
36
    scannet_point_cloud = scannet_results.get('points', None)
liyinhao's avatar
liyinhao committed
37
38
39
40
41
42
43
44
45
46
    assert scannet_point_cloud.shape == (1000, 4)


def test_load_annotations3D():
    sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
    sunrgbd_load_annotations3D = LoadAnnotations3D()
    sunrgbd_results = dict()
    sunrgbd_results['data_path'] = './tests/data/sunrgbd/sunrgbd_trainval'
    sunrgbd_results['info'] = sunrgbd_info[0]
    sunrgbd_results = sunrgbd_load_annotations3D(sunrgbd_results)
liyinhao's avatar
liyinhao committed
47
    sunrgbd_gt_boxes = sunrgbd_results.get('gt_bboxes_3d', None)
liyinhao's avatar
liyinhao committed
48
    sunrgbd_gt_lbaels = sunrgbd_results.get('gt_labels', None)
liyinhao's avatar
liyinhao committed
49
    sunrgbd_gt_boxes_mask = sunrgbd_results.get('gt_bboxes_3d_mask', None)
liyinhao's avatar
liyinhao committed
50
    assert sunrgbd_gt_boxes.shape == (3, 7)
liyinhao's avatar
liyinhao committed
51
    assert sunrgbd_gt_lbaels.shape == (3, 1)
liyinhao's avatar
liyinhao committed
52
    assert sunrgbd_gt_boxes_mask.shape == (3, 1)
liyinhao's avatar
liyinhao committed
53

54
    scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
liyinhao's avatar
liyinhao committed
55
    scannet_load_annotations3D = LoadAnnotations3D()
liyinhao's avatar
liyinhao committed
56
    scannet_results = dict()
liyinhao's avatar
liyinhao committed
57
58
59
60
61
62
63
64
65
    data_path = './tests/data/scannet/scannet_train_instance_data'
    scannet_results['data_path'] = data_path
    scannet_info = scannet_info[0]
    scan_name = scannet_info['point_cloud']['lidar_idx']
    scannet_results['ins_labelname'] = osp.join(data_path,
                                                scan_name + '_ins_label.npy')
    scannet_results['sem_labelname'] = osp.join(data_path,
                                                scan_name + '_sem_label.npy')
    scannet_results['info'] = scannet_info
liyinhao's avatar
liyinhao committed
66
    scannet_results = scannet_load_annotations3D(scannet_results)
liyinhao's avatar
liyinhao committed
67
    scannet_gt_boxes = scannet_results.get('gt_bboxes_3d', None)
liyinhao's avatar
liyinhao committed
68
    scannet_gt_lbaels = scannet_results.get('gt_labels', None)
liyinhao's avatar
liyinhao committed
69
    scannet_gt_boxes_mask = scannet_results.get('gt_bboxes_3d_mask', None)
liyinhao's avatar
liyinhao committed
70
71
    scannet_pts_instance_mask = scannet_results.get('pts_instance_mask', None)
    scannet_pts_semantic_mask = scannet_results.get('pts_semantic_mask', None)
liyinhao's avatar
liyinhao committed
72
    assert scannet_gt_boxes.shape == (27, 6)
liyinhao's avatar
liyinhao committed
73
    assert scannet_gt_lbaels.shape == (27, 1)
liyinhao's avatar
liyinhao committed
74
    assert scannet_gt_boxes_mask.shape == (27, 1)
liyinhao's avatar
liyinhao committed
75
76
    assert scannet_pts_instance_mask.shape == (1000, )
    assert scannet_pts_semantic_mask.shape == (1000, )