viz_bbox_predictions.py

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()