File size: 5,326 Bytes
3403979 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | import utils
import logging
import argparse
import importlib
import torch
import numpy as np
import matplotlib.pyplot as plt
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 models.utils import DUMP, VERSION
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)
parser.add_argument('--stage_id', default=5)
parser.add_argument('--num_frames', default=3)
parser.add_argument('--num_views', default=6)
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=2,
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']
for idx, data in enumerate(val_loader):
DUMP.enabled = True
model.eval()
with torch.no_grad():
model(return_loss=False, rescale=True, **data)
cls_scores = torch.load('{}/cls_score_stage{}.pth'.format(DUMP.out_dir, args.stage_id))[0]
cls_scores, cls_ids = torch.max(cls_scores, dim=-1)
# only select queries with high confidence
query_ids = torch.where(cls_scores > args.score_threshold)[0]
cls_scores, cls_ids = cls_scores[query_ids], cls_ids[query_ids]
plt.figure(figsize=(240, 49))
view_mapping = [1, 2, 0, 4, 5, 3]
for frame_id in range(args.num_frames):
sample_points_cam = torch.load(
'{}/sample_points_cam_stage{}.pth'.format(DUMP.out_dir, args.stage_id)
) # [1, 8f, 6view, 900, 32, 2]
valid_mask = torch.load(
'{}/sample_points_cam_valid_mask_stage{}.pth'.format(DUMP.out_dir, args.stage_id)
) # [1, 8f, 6view, 900, 32]
for view_id in range(args.num_views):
filenames = data['img_metas'][0].data[0][0]['filename']
filename = filenames[frame_id * 6 + view_id]
# crop 1600x640 area
img = Image.open(filename)
img = img.crop((0, 260, 1600, 900))
# plot image
plot_id = frame_id * args.num_views + view_mapping[view_id] + 1
ax = plt.subplot(args.num_frames, args.num_views, plot_id)
ax.imshow(img)
ax.axis('off')
ax.set_xlim(0, 1600)
ax.set_ylim(640, 0)
# plot the sampling points for each query
for query_id in query_ids:
xyz = sample_points_cam[0, frame_id, view_id, query_id].numpy() # [32, 3]
mask = valid_mask[0, frame_id, view_id, query_id].numpy() # [32]
mask = np.round(mask).astype(bool)
cx = xyz[:, 0] * 1600
cy = xyz[:, 1] * 640
cz = xyz[:, 2]
cz[np.where(cz <= 0)] = 1e8
cz = np.log(60 / cz ** 0.8) * 2.4
cx, cy, cz = cx[mask], cy[mask], cz[mask]
if len(cz) == 0:
continue
ax.scatter(cx, cy, s=4**(cz + 1), alpha=0.7, color='C%d' % (query_id % 5))
plt.tight_layout()
plt.subplots_adjust(hspace=0.01, wspace=0.01)
plt.savefig('outputs/sp_%04d.jpg' % idx, dpi=20)
plt.close()
logging.info('Visualized result is dumped to outputs/sp_%04d.jpg' % idx)
if __name__ == '__main__':
main()
|