MonoCon/mmdetection3d-0.14.0/tests/test_data/test_datasets/test_waymo_dataset.py

262 lines
10 KiB
Python

import numpy as np
import pytest
import tempfile
import torch
from mmdet3d.datasets import WaymoDataset
def _generate_waymo_train_dataset_config():
data_root = 'tests/data/waymo/kitti_format/'
ann_file = 'tests/data/waymo/kitti_format/waymo_infos_train.pkl'
classes = ['Car', 'Pedestrian', 'Cyclist']
pts_prefix = 'velodyne'
point_cloud_range = [-74.88, -74.88, -2, 74.88, 74.88, 4]
file_client_args = dict(backend='disk')
db_sampler = dict(
data_root=data_root,
info_path=data_root + 'waymo_dbinfos_train.pkl',
rate=1.0,
prepare=dict(
filter_by_difficulty=[-1], filter_by_min_points=dict(Car=5)),
classes=classes,
sample_groups=dict(Car=15),
points_loader=dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=5,
use_dim=[0, 1, 2, 3, 4],
file_client_args=file_client_args))
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=6,
use_dim=5,
file_client_args=file_client_args),
dict(
type='LoadAnnotations3D',
with_bbox_3d=True,
with_label_3d=True,
file_client_args=file_client_args),
dict(type='ObjectSample', db_sampler=db_sampler),
dict(
type='RandomFlip3D',
sync_2d=False,
flip_ratio_bev_horizontal=0.5,
flip_ratio_bev_vertical=0.5),
dict(
type='GlobalRotScaleTrans',
rot_range=[-0.78539816, 0.78539816],
scale_ratio_range=[0.95, 1.05]),
dict(type='PointsRangeFilter', point_cloud_range=point_cloud_range),
dict(type='ObjectRangeFilter', point_cloud_range=point_cloud_range),
dict(type='PointShuffle'),
dict(type='DefaultFormatBundle3D', class_names=classes),
dict(
type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels_3d'])
]
modality = dict(use_lidar=True, use_camera=False)
split = 'training'
return data_root, ann_file, classes, pts_prefix, pipeline, modality, split
def _generate_waymo_val_dataset_config():
data_root = 'tests/data/waymo/kitti_format/'
ann_file = 'tests/data/waymo/kitti_format/waymo_infos_val.pkl'
classes = ['Car', 'Pedestrian', 'Cyclist']
pts_prefix = 'velodyne'
point_cloud_range = [-74.88, -74.88, -2, 74.88, 74.88, 4]
file_client_args = dict(backend='disk')
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=6,
use_dim=5,
file_client_args=file_client_args),
dict(
type='MultiScaleFlipAug3D',
img_scale=(1333, 800),
pts_scale_ratio=1,
flip=False,
transforms=[
dict(
type='GlobalRotScaleTrans',
rot_range=[0, 0],
scale_ratio_range=[1., 1.],
translation_std=[0, 0, 0]),
dict(type='RandomFlip3D'),
dict(
type='PointsRangeFilter',
point_cloud_range=point_cloud_range),
dict(
type='DefaultFormatBundle3D',
class_names=classes,
with_label=False),
dict(type='Collect3D', keys=['points'])
])
]
modality = dict(use_lidar=True, use_camera=False)
split = 'training'
return data_root, ann_file, classes, pts_prefix, pipeline, modality, split
def test_getitem():
np.random.seed(0)
data_root, ann_file, classes, pts_prefix, pipeline, \
modality, split = _generate_waymo_train_dataset_config()
waymo_dataset = WaymoDataset(data_root, ann_file, split, pts_prefix,
pipeline, classes, modality)
data = waymo_dataset[0]
points = data['points']._data
gt_bboxes_3d = data['gt_bboxes_3d']._data
gt_labels_3d = data['gt_labels_3d']._data
expected_gt_bboxes_3d = torch.tensor(
[[31.4750, -4.5690, 2.1857, 2.3519, 6.0931, 3.1756, -1.2895]])
expected_gt_labels_3d = torch.tensor([0])
assert points.shape == (765, 5)
assert torch.allclose(
gt_bboxes_3d.tensor, expected_gt_bboxes_3d, atol=1e-4)
assert torch.all(gt_labels_3d == expected_gt_labels_3d)
def test_evaluate():
if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda')
from mmdet3d.core.bbox import LiDARInstance3DBoxes
data_root, ann_file, classes, pts_prefix, pipeline, \
modality, split = _generate_waymo_val_dataset_config()
waymo_dataset = WaymoDataset(data_root, ann_file, split, pts_prefix,
pipeline, classes, modality)
boxes_3d = LiDARInstance3DBoxes(
torch.tensor([[
6.9684e+01, 3.3335e+01, 4.1465e-02, 2.0100e+00, 4.3600e+00,
1.4600e+00, -9.0000e-02
]]))
labels_3d = torch.tensor([0])
scores_3d = torch.tensor([0.5])
result = dict(boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d)
# kitti protocol
metric = ['kitti']
ap_dict = waymo_dataset.evaluate([result], metric=metric)
assert np.isclose(ap_dict['KITTI/Overall_3D_easy'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_moderate'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_hard'], 3.0303030303030307)
# waymo protocol
metric = ['waymo']
boxes_3d = LiDARInstance3DBoxes(
torch.tensor([[
6.9684e+01, 3.3335e+01, 4.1465e-02, 2.0100e+00, 4.3600e+00,
1.4600e+00, -9.0000e-02
]]))
labels_3d = torch.tensor([0])
scores_3d = torch.tensor([0.8])
result = dict(boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d)
ap_dict = waymo_dataset.evaluate([result], metric=metric)
assert np.isclose(ap_dict['Overall/L1 mAP'], 0.3333333333333333)
assert np.isclose(ap_dict['Overall/L2 mAP'], 0.3333333333333333)
assert np.isclose(ap_dict['Overall/L1 mAPH'], 0.3333333333333333)
assert np.isclose(ap_dict['Overall/L2 mAPH'], 0.3333333333333333)
def test_show():
import mmcv
from os import path as osp
from mmdet3d.core.bbox import LiDARInstance3DBoxes
# Waymo shares show function with KITTI so I just copy it here
tmp_dir = tempfile.TemporaryDirectory()
temp_dir = tmp_dir.name
data_root, ann_file, classes, pts_prefix, pipeline, \
modality, split = _generate_waymo_val_dataset_config()
waymo_dataset = WaymoDataset(
data_root, ann_file, split=split, modality=modality, pipeline=pipeline)
boxes_3d = LiDARInstance3DBoxes(
torch.tensor(
[[46.1218, -4.6496, -0.9275, 0.5316, 1.4442, 1.7450, 1.1749],
[33.3189, 0.1981, 0.3136, 0.5656, 1.2301, 1.7985, 1.5723],
[46.1366, -4.6404, -0.9510, 0.5162, 1.6501, 1.7540, 1.3778],
[33.2646, 0.2297, 0.3446, 0.5746, 1.3365, 1.7947, 1.5430],
[58.9079, 16.6272, -1.5829, 1.5656, 3.9313, 1.4899, 1.5505]]))
scores_3d = torch.tensor([0.1815, 0.1663, 0.5792, 0.2194, 0.2780])
labels_3d = torch.tensor([0, 0, 1, 1, 2])
result = dict(boxes_3d=boxes_3d, scores_3d=scores_3d, labels_3d=labels_3d)
results = [result]
waymo_dataset.show(results, temp_dir, show=False)
pts_file_path = osp.join(temp_dir, '1000000', '1000000_points.obj')
gt_file_path = osp.join(temp_dir, '1000000', '1000000_gt.obj')
pred_file_path = osp.join(temp_dir, '1000000', '1000000_pred.obj')
mmcv.check_file_exist(pts_file_path)
mmcv.check_file_exist(gt_file_path)
mmcv.check_file_exist(pred_file_path)
tmp_dir.cleanup()
# test show with pipeline
tmp_dir = tempfile.TemporaryDirectory()
temp_dir = tmp_dir.name
eval_pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=6,
use_dim=5),
dict(
type='DefaultFormatBundle3D',
class_names=classes,
with_label=False),
dict(type='Collect3D', keys=['points'])
]
waymo_dataset.show(results, temp_dir, show=False, pipeline=eval_pipeline)
pts_file_path = osp.join(temp_dir, '1000000', '1000000_points.obj')
gt_file_path = osp.join(temp_dir, '1000000', '1000000_gt.obj')
pred_file_path = osp.join(temp_dir, '1000000', '1000000_pred.obj')
mmcv.check_file_exist(pts_file_path)
mmcv.check_file_exist(gt_file_path)
mmcv.check_file_exist(pred_file_path)
tmp_dir.cleanup()
def test_format_results():
if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda')
from mmdet3d.core.bbox import LiDARInstance3DBoxes
data_root, ann_file, classes, pts_prefix, pipeline, \
modality, split = _generate_waymo_val_dataset_config()
waymo_dataset = WaymoDataset(data_root, ann_file, split, pts_prefix,
pipeline, classes, modality)
boxes_3d = LiDARInstance3DBoxes(
torch.tensor([[
6.9684e+01, 3.3335e+01, 4.1465e-02, 2.0100e+00, 4.3600e+00,
1.4600e+00, -9.0000e-02
]]))
labels_3d = torch.tensor([0])
scores_3d = torch.tensor([0.5])
result = dict(boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d)
result_files, tmp_dir = waymo_dataset.format_results([result],
data_format='waymo')
expected_name = np.array(['Car'])
expected_truncated = np.array([0.])
expected_occluded = np.array([0])
expected_alpha = np.array([0.35619745])
expected_bbox = np.array([[0., 673.59814, 37.07779, 719.7537]])
expected_dimensions = np.array([[4.36, 1.46, 2.01]])
expected_location = np.array([[-33.000042, 2.4999967, 68.29972]])
expected_rotation_y = np.array([-0.09])
expected_score = np.array([0.5])
expected_sample_idx = np.array([1000000])
assert np.all(result_files[0]['name'] == expected_name)
assert np.allclose(result_files[0]['truncated'], expected_truncated)
assert np.all(result_files[0]['occluded'] == expected_occluded)
assert np.allclose(result_files[0]['alpha'], expected_alpha)
assert np.allclose(result_files[0]['bbox'], expected_bbox)
assert np.allclose(result_files[0]['dimensions'], expected_dimensions)
assert np.allclose(result_files[0]['location'], expected_location)
assert np.allclose(result_files[0]['rotation_y'], expected_rotation_y)
assert np.allclose(result_files[0]['score'], expected_score)
assert np.allclose(result_files[0]['sample_idx'], expected_sample_idx)
tmp_dir.cleanup()