-
Notifications
You must be signed in to change notification settings - Fork 27
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
78bb278
commit e9dca31
Showing
6 changed files
with
420 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,244 @@ | ||
| import utils | ||
| import logging | ||
| import argparse | ||
| import importlib | ||
| import torch | ||
| import numpy as np | ||
| import matplotlib.pyplot as plt | ||
| from matplotlib.gridspec import GridSpec | ||
| from PIL import Image | ||
| from mmcv import Config, DictAction | ||
| from mmcv.parallel import MMDataParallel | ||
| from mmcv.runner import load_checkpoint | ||
| from mmdet.apis import set_random_seed | ||
| from mmdet3d.datasets import build_dataset, build_dataloader | ||
| from mmdet3d.models import build_model | ||
| from nuscenes.utils.data_classes import Box | ||
| from pyquaternion import Quaternion | ||
| from nuscenes.nuscenes import NuScenes | ||
| from nuscenes.utils.geometry_utils import box_in_image | ||
| from configs.r50_nuimg_704x256 import class_names | ||
| from models.utils import VERSION | ||
|
|
||
|
|
||
| classname_to_color = { # RGB | ||
| 'car': (255, 158, 0), # Orange | ||
| 'pedestrian': (0, 0, 230), # Blue | ||
| 'trailer': (255, 140, 0), # Darkorange | ||
| 'truck': (255, 99, 71), # Tomato | ||
| 'bus': (255, 127, 80), # Coral | ||
| 'motorcycle': (255, 61, 99), # Red | ||
| 'construction_vehicle': (233, 150, 70), # Darksalmon | ||
| 'bicycle': (220, 20, 60), # Crimson | ||
| 'barrier': (112, 128, 144), # Slategrey | ||
| 'traffic_cone': (47, 79, 79), # Darkslategrey | ||
| } | ||
|
|
||
|
|
||
| def convert_to_nusc_box(bboxes, scores=None, labels=None, names=None, score_threshold=0.3, lift_center=False): | ||
| results = [] | ||
| for q in range(bboxes.shape[0]): | ||
| if scores is not None: | ||
| score = scores[q] | ||
| else: | ||
| score = 1.0 | ||
|
|
||
| if score < score_threshold: | ||
| continue | ||
|
|
||
| if labels is not None: | ||
| label = labels[q] | ||
| else: | ||
| label = 0 | ||
|
|
||
| if names is not None: | ||
| name = names[q] | ||
| else: | ||
| name = class_names[label] | ||
|
|
||
| if name not in class_names: | ||
| name = class_names[-1] | ||
|
|
||
| bbox = bboxes[q].copy() | ||
| if lift_center: | ||
| bbox[2] += bbox[5] * 0.5 | ||
|
|
||
| orientation = Quaternion(axis=[0, 0, 1], radians=bbox[6]) | ||
|
|
||
| box = Box( | ||
| center=[bbox[0], bbox[1], bbox[2]], | ||
| size=[bbox[4], bbox[3], bbox[5]], | ||
| orientation=orientation, | ||
| score=score, | ||
| label=label, | ||
| velocity=(bbox[7], bbox[8], 0), | ||
| name=name | ||
| ) | ||
|
|
||
| results.append(box) | ||
|
|
||
| return results | ||
|
|
||
|
|
||
| def viz_bbox(nusc, bboxes, data_info, fig, gs): | ||
| cam_types = [ | ||
| 'CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', | ||
| 'CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT', | ||
| ] | ||
|
|
||
| for cam_id, cam_type in enumerate(cam_types): | ||
| sample_data_token = nusc.get('sample', data_info['token'])['data'][cam_type] | ||
|
|
||
| sd_record = nusc.get('sample_data', sample_data_token) | ||
| cs_record = nusc.get('calibrated_sensor', sd_record['calibrated_sensor_token']) | ||
| intrinsic = np.array(cs_record['camera_intrinsic']) | ||
|
|
||
| img_path = nusc.get_sample_data_path(sample_data_token) | ||
| img_size = (sd_record['width'], sd_record['height']) | ||
|
|
||
| ax = fig.add_subplot(gs[cam_id // 3, cam_id % 3]) | ||
| ax.imshow(Image.open(img_path)) | ||
|
|
||
| for bbox in bboxes: | ||
| bbox = bbox.copy() | ||
|
|
||
| # Move box to ego vehicle coord system | ||
| bbox.rotate(Quaternion(data_info['lidar2ego_rotation'])) | ||
| bbox.translate(np.array(data_info['lidar2ego_translation'])) | ||
|
|
||
| # Move box to sensor coord system | ||
| bbox.translate(-np.array(cs_record['translation'])) | ||
| bbox.rotate(Quaternion(cs_record['rotation']).inverse) | ||
|
|
||
| if box_in_image(bbox, intrinsic, img_size): | ||
| c = np.array(classname_to_color[bbox.name]) / 255.0 | ||
| bbox.render(ax, view=intrinsic, normalize=True, colors=(c, c, c), linewidth=1) | ||
|
|
||
| ax.axis('off') | ||
| ax.set_title(cam_type) | ||
| ax.set_xlim(0, img_size[0]) | ||
| ax.set_ylim(img_size[1], 0) | ||
|
|
||
| sample = nusc.get('sample', data_info['token']) | ||
| lidar_data_token = sample['data']['LIDAR_TOP'] | ||
|
|
||
| ax = fig.add_subplot(gs[0:2, 3]) | ||
| nusc.explorer.render_sample_data(lidar_data_token, with_anns=False, ax=ax, verbose=False) | ||
| ax.axis('off') | ||
| ax.set_title('LIDAR_TOP') | ||
| ax.set_xlim(-40, 40) | ||
| ax.set_ylim(-40, 40) | ||
|
|
||
| sd_record = nusc.get('sample_data', lidar_data_token) | ||
| pose_record = nusc.get('ego_pose', sd_record['ego_pose_token']) | ||
| cs_record = nusc.get('calibrated_sensor', sd_record['calibrated_sensor_token']) | ||
|
|
||
| for bbox in bboxes: | ||
| bbox = bbox.copy() | ||
|
|
||
| bbox.rotate(Quaternion(cs_record['rotation'])) | ||
| bbox.translate(np.array(cs_record['translation'])) | ||
| bbox.rotate(Quaternion(pose_record['rotation'])) | ||
|
|
||
| yaw = Quaternion(pose_record['rotation']).yaw_pitch_roll[0] | ||
| bbox.rotate(Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]).inverse) | ||
|
|
||
| c = np.array(classname_to_color[bbox.name]) / 255.0 | ||
| bbox.render(ax, view=np.eye(4), colors=(c, c, c)) | ||
|
|
||
|
|
||
| def main(): | ||
| parser = argparse.ArgumentParser(description='Validate a detector') | ||
| parser.add_argument('--config', required=True) | ||
| parser.add_argument('--weights', required=True) | ||
| parser.add_argument('--override', nargs='+', action=DictAction) | ||
| parser.add_argument('--score_threshold', default=0.3) | ||
| args = parser.parse_args() | ||
|
|
||
| # parse configs | ||
| cfgs = Config.fromfile(args.config) | ||
| if args.override is not None: | ||
| cfgs.merge_from_dict(args.override) | ||
|
|
||
| # use val-mini for visualization | ||
| cfgs.data.val.ann_file = cfgs.data.val.ann_file.replace('val', 'val_mini') | ||
|
|
||
| # register custom module | ||
| importlib.import_module('models') | ||
| importlib.import_module('loaders') | ||
|
|
||
| # MMCV, please shut up | ||
| from mmcv.utils.logging import logger_initialized | ||
| logger_initialized['root'] = logging.Logger(__name__, logging.WARNING) | ||
| logger_initialized['mmcv'] = logging.Logger(__name__, logging.WARNING) | ||
|
|
||
| # you need one GPU | ||
| assert torch.cuda.is_available() | ||
| assert torch.cuda.device_count() == 1 | ||
|
|
||
| utils.init_logging(None, cfgs.debug) | ||
| logging.info('Using GPU: %s' % torch.cuda.get_device_name(0)) | ||
| logging.info('Setting random seed: 0') | ||
| set_random_seed(0, deterministic=True) | ||
|
|
||
| logging.info('Loading validation set from %s' % cfgs.data.val.data_root) | ||
| val_dataset = build_dataset(cfgs.data.val) | ||
| val_loader = build_dataloader( | ||
| val_dataset, | ||
| samples_per_gpu=1, | ||
| workers_per_gpu=cfgs.data.workers_per_gpu, | ||
| num_gpus=1, | ||
| dist=False, | ||
| shuffle=False, | ||
| seed=0, | ||
| ) | ||
|
|
||
| logging.info('Creating model: %s' % cfgs.model.type) | ||
| model = build_model(cfgs.model) | ||
| model.cuda() | ||
| model = MMDataParallel(model, [0]) | ||
|
|
||
| logging.info('Loading checkpoint from %s' % args.weights) | ||
| checkpoint = load_checkpoint( | ||
| model, args.weights, map_location='cuda', strict=True, | ||
| logger=logging.Logger(__name__, logging.ERROR) | ||
| ) | ||
|
|
||
| if 'version' in checkpoint: | ||
| VERSION.name = checkpoint['version'] | ||
|
|
||
| logging.info('Initialize nuscenes toolkit...') | ||
| if 'mini' in cfgs.data.val.ann_file: | ||
| nusc = NuScenes(version='v1.0-mini', dataroot=cfgs.data.val.data_root, verbose=False) | ||
| else: | ||
| nusc = NuScenes(version='v1.0-trainval', dataroot=cfgs.data.val.data_root, verbose=False) | ||
|
|
||
| for i, data in enumerate(val_loader): | ||
| model.eval() | ||
|
|
||
| with torch.no_grad(): | ||
| results = model(return_loss=False, rescale=True, **data) | ||
| results = results[0]['pts_bbox'] | ||
|
|
||
| bboxes_pred = convert_to_nusc_box( | ||
| bboxes=results['boxes_3d'].tensor.numpy(), | ||
| scores=results['scores_3d'].numpy(), | ||
| labels=results['labels_3d'].numpy(), | ||
| score_threshold=args.score_threshold, | ||
| lift_center=True, | ||
| ) | ||
|
|
||
| fig = plt.figure(figsize=(15.5, 5)) | ||
| gs = GridSpec(2, 4, figure=fig) | ||
|
|
||
| viz_bbox(nusc, bboxes_pred, val_dataset.data_infos[i], fig, gs) | ||
|
|
||
| plt.tight_layout() | ||
| plt.savefig('outputs/bbox_%04d.jpg' % i, dpi=200) | ||
| plt.close() | ||
|
|
||
| logging.info('Visualized result is dumped to outputs/bbox_%04d.jpg' % i) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| main() |
Oops, something went wrong.