Unverified Commit 1e7059ff authored by xiliu8006's avatar xiliu8006 Committed by GitHub
Browse files

[Feature] Support nuimages visualization (#408)

* support nuimages visualization

* recover config

* support nuimages visualization

* support batchsize > 1

* add comments
parent 9cb75e7d
import mmcv
import os
import torch
from mmcv.image import tensor2imgs
def single_gpu_test(model, data_loader, show=False, out_dir=None):
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
"""Test model with single gpu.
This method tests model with single gpu and gives the 'show' option.
......@@ -29,8 +35,42 @@ def single_gpu_test(model, data_loader, show=False, out_dir=None):
result = model(return_loss=False, rescale=True, **data)
if show:
model.module.show_results(data, result, out_dir)
# Visualize the results of MMdetection3D model
# 'show_results' is MMdetection3D visualization API
if hasattr(model.module, 'show_results'):
model.module.show_results(data, result, out_dir)
# Visualize the results of MMdetection model
# 'show_result' is MMdetection visualization API
else:
batch_size = len(result)
if batch_size == 1 and isinstance(data['img'][0],
torch.Tensor):
img_tensor = data['img'][0]
else:
img_tensor = data['img'][0].data[0]
img_metas = data['img_metas'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
assert len(imgs) == len(img_metas)
for i, (img, img_meta) in enumerate(zip(imgs, img_metas)):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
ori_h, ori_w = img_meta['ori_shape'][:-1]
img_show = mmcv.imresize(img_show, (ori_w, ori_h))
if out_dir:
out_file = os.path.join(out_dir,
img_meta['ori_filename'])
else:
out_file = None
model.module.show_result(
img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
results.extend(result)
batch_size = len(result)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment