def multi_gpu_test(model, dataset, cfg, show=False, tmpdir=None):
model.eval()
results = []
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for idx in range(rank, len(dataset), world_size):
data = dataset[idx]
# None type data cannot be scatter, here we pick out the not None type data
notNoneData = {}
for k, v in zip(data.keys(), data.values()):
if v is not None:
notNoneData[k] = v
notNoneData = scatter(
collate([notNoneData], samples_per_gpu=1),
[torch.cuda.current_device()]
)[0]
data.update(notNoneData)
with torch.no_grad():
result = model(data)
filter_result = {}
filter_result.update(Error=result['Error'])
if show:
item = dataset.data_list[idx]
result['leftImage'] = imread(
osp.join(cfg.data.test.data_root, item['left_image_path'])
).astype(np.float32)
result['rightImage'] = imread(
osp.join(cfg.data.test.data_root, item['right_image_path'])
).astype(np.float32)
image_name = item['left_image_path'].split('/')[-1]
save_result(result, cfg.out_dir, image_name)
if hasattr(cfg, 'sparsification_plot'):
filter_result['Error'].update(sparsification_eval(result, cfg))
results.append(filter_result)
if rank == 0:
batch_size = world_size
for _ in range(batch_size):
prog_bar.update()
# collect results from all ranks
results = collect_results(results, len(dataset), tmpdir)
return results
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
if self.task == 'BddStreet':
lane_bins = [np.zeros((n, n)) for n in [3, 3, 9]]
drivable_bin = np.zeros((3, 3))
elif self.task == 'BddSemanticSeg':
sem_seg_bin = np.zeros((20, 20))
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in range(len(self.dataset)):
data = self.dataset[idx]
data_gpu = scatter(
collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(
return_loss=False, rescale=True, **data_gpu)
# evaluation
if self.task == 'BddStreet':
if self.dataset.with_lane:
_lane_bins = self.eval_lane(result['lane_results'], data['img_meta'][0].data['gt_lane'])
for i, bin in enumerate(_lane_bins):
lane_bins[i] += bin
if self.dataset.with_drivable:
drivable_bin += self.eval_drivable(result['drivable_results'], data['img_meta'][0].data['gt_drivable'])
elif self.task == 'BddSemanticSeg':
sem_seg_bin += self.eval_sem_seg(result['sem_seg_results'], data['img_meta'][0].data['gt_sem_seg'])
prog_bar.update()
# lane_IoUs = [self.fg_IoUs(b) for b in lane_bins]
# drivable_IoUs = self.fg_IoUs(drivable_bin)
outputs = dict()
if self.task == 'BddStreet':
print('\nSTREET EVALUATION')
if self.dataset.with_lane:
lane_mIoUs = [self.fg_mIoU(bin) for bin in lane_bins]
print('[lane] direction: {} style: {} type: {}'.format(*lane_mIoUs))
runner.log_buffer.output.update(dict(val_lane_0=lane_mIoUs[0], val_lane_1=lane_mIoUs[1], val_lane_2=lane_mIoUs[2]))
lane_avg_recall = [self.fg_avg_recall(bin) for bin in lane_bins]
print('[lane] direction: {} style: {} type: {}'.format(*lane_avg_recall))
runner.log_buffer.output.update(dict(val_lane_ar_0=lane_avg_recall[0], val_lane_ar_1=lane_avg_recall[1], val_lane_ar_2=lane_avg_recall[2]))
if self.dataset.with_drivable:
drivable_mIoU = self.fg_mIoU(drivable_bin)
print('[driv] mIoU: {}'.format(drivable_mIoU))
runner.log_buffer.output.update(dict(val_drivable=drivable_mIoU))
elif self.task == 'BddSemanticSeg':
sem_seg_mIoU = self.fg_mIoU(sem_seg_bin)
print('\nSEMANTIC SEG EVALUATION\n[sem_seg] mIoU: {}'.format(sem_seg_mIoU))
runner.log_buffer.output.update(dict(val_sem_seg=sem_seg_mIoU))
runner.log_buffer.ready = True
def generation_inference(model, img, img_unpaired=None):
"""Inference image with the model.
Args:
model (nn.Module): The loaded model.
img (str): File path of input image.
img_unpaired (str, optional): File path of the unpaired image.
If not None, perform unpaired image generation. Default: None.
Returns:
np.ndarray: The predicted generation result.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
# prepare data
if img_unpaired is None:
data = dict(pair_path=img)
else:
data = dict(img_a_path=img, img_b_path=img_unpaired)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# forward the model
with torch.no_grad():
results = model(test_mode=True, **data)
# process generation shown mode
if img_unpaired is None:
if model.show_input:
output = np.concatenate([
tensor2img(results['real_a'], min_max=(-1, 1)),
tensor2img(results['fake_b'], min_max=(-1, 1)),
tensor2img(results['real_b'], min_max=(-1, 1))
],
axis=1)
else:
output = tensor2img(results['fake_b'], min_max=(-1, 1))
else:
if model.show_input:
output = np.concatenate([
tensor2img(results['real_a'], min_max=(-1, 1)),
tensor2img(results['fake_b'], min_max=(-1, 1)),
tensor2img(results['real_b'], min_max=(-1, 1)),
tensor2img(results['fake_a'], min_max=(-1, 1))
],
axis=1)
else:
if model.test_direction == 'a2b':
output = tensor2img(results['fake_b'], min_max=(-1, 1))
else:
output = tensor2img(results['fake_a'], min_max=(-1, 1))
return output
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
range_idxs = list(
range(runner.rank, len(self.dataset), runner.world_size))
if self.shuffle:
np.random.shuffle(range_idxs)
range_idxs = range_idxs[:self.num_evals]
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(range_idxs) * runner.world_size)
results = []
for idx in range_idxs:
data = self.dataset[idx]
data_gpu = scatter(collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(return_loss=False,
rescale=True,
**data_gpu)
results.append(result)
batch_size = runner.world_size
if runner.rank == 0:
for _ in range(batch_size):
prog_bar.update()
if runner.rank == 0:
print('\n')
dist.barrier()
for i in range(1, runner.world_size):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_range = osp.join(runner.work_dir,
'temp_range_{}.pkl'.format(i))
tmp_results = mmcv.load(tmp_file)
tmp_range_idxs = mmcv.load(tmp_range)
results.extend(tmp_results)
range_idxs.extend(tmp_range_idxs)
os.remove(tmp_file)
os.remove(tmp_range)
self.evaluate(runner, results, range_idxs)
else:
tmp_file = osp.join(runner.work_dir,
'temp_{}.pkl'.format(runner.rank))
tmp_range = osp.join(runner.work_dir,
'temp_range_{}.pkl'.format(runner.rank))
mmcv.dump(results, tmp_file)
mmcv.dump(range_idxs, tmp_range)
dist.barrier()
dist.barrier()
def matting_inference(model, img, trimap):
"""Inference image(s) with the model.
Args:
model (nn.Module): The loaded model.
img (str): Image file path.
trimap (str): Trimap file path.
Returns:
np.ndarray: The predicted alpha matte.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# remove alpha from test_pipeline
keys_to_remove = ['alpha', 'ori_alpha']
for key in keys_to_remove:
for pipeline in list(cfg.test_pipeline):
if 'key' in pipeline and key == pipeline['key']:
cfg.test_pipeline.remove(pipeline)
if 'keys' in pipeline and key in pipeline['keys']:
pipeline['keys'].remove(key)
if len(pipeline['keys']) == 0:
cfg.test_pipeline.remove(pipeline)
if 'meta_keys' in pipeline and key in pipeline['meta_keys']:
pipeline['meta_keys'].remove(key)
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
# prepare data
data = dict(merged_path=img, trimap_path=trimap)
data = test_pipeline(data)
# # Test Code
# merged = data['merged']
# ori_merged = data['ori_merged']
# trimap = data['trimap']
# trimap_transformed = data['trimap_transformed']
# ori_merged.cpu().numpy().tofile('dat/' + 'ori_merged_new' + '.dat')
# merged.cpu().numpy().tofile('dat/' + 'merged_rgbtrue' + '.dat')
# trimap.cpu().numpy().tofile('dat/' + 'trimap' + '.dat')
# trimap_transformed.numpy().tofile('dat/' + 'trimap_transformed' + '.dat')
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# merged = data['merged']
# merged.cpu().numpy().tofile('dat/merged_tensor_new_norm_list.dat')
print("data prepare success!!!")
# forward the model
with torch.no_grad():
result = model(test_mode=True, **data)
return result['pred_alpha']
def load_image(self, img):
cfg = self.model.cfg
device = next(self.model.parameters()).device # model device
# prepare data
data = dict(img_info=dict(filename=img['img_path']), img_prefix=None)
# build the data pipeline
test_pipeline = Compose(cfg.data.test.pipeline)
# we store the metadata about mmdetection's pipeline transformations in data. It will be used later
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
# scatter to specified GPU
self.preprocessed_images = scatter(data, [device])[0]
self.features = self.model.extract_feats(self.preprocessed_images['img'])[0]
def inference_detector(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# Use torchvision ops for CPU mode instead
for m in model.modules():
if isinstance(m, (RoIPool, RoIAlign)):
if not m.aligned:
# aligned=False is not implemented on CPU
# set use_torchvision on-the-fly
m.use_torchvision = True
warnings.warn('We set use_torchvision=True in CPU mode.')
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
# forward the model
with torch.no_grad():
data['img'][0] = data['img'][0].cuda()
data.update({'x': data.pop('img')})
result = model(return_loss=False, rescale=True, **data)
y_head_f_1, y_head_f_2, y_head_cls = model(return_loss=False, rescale=True, return_box=False, **data)
y_head_f_1 = torch.cat(y_head_f_1, 0)
y_head_f_2 = torch.cat(y_head_f_2, 0)
y_head_f_1 = torch.nn.Sigmoid()(y_head_f_1)
y_head_f_2 = torch.nn.Sigmoid()(y_head_f_2)
loss_l2_p = (y_head_f_1 - y_head_f_2).pow(2)
uncertainty_all_N = loss_l2_p.mean(dim=1)
arg = uncertainty_all_N.argsort()
uncertainty_single = uncertainty_all_N[arg[-cfg.k:]].mean()
return result, uncertainty_single
def inference_detector_(model, img, img_pair, temp_feats=None):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
# t0 = time.time()
data1 = dict(img=img)
data1 = test_pipeline(data1)
if temp_feats is None:
data2 = dict(img=img_pair)
data2 = test_pipeline(data2)
data2_img = data2['img']
for i in range(len(data1['img'])):
data1['img'][i] = torch.cat((data1['img'][i].data, data2_img[i].data))
data = data1
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
with torch.no_grad():
result, temp_feats = model(return_loss=False, temp_feats=None, rescale=True, **data)
return result, temp_feats
else:
data = data1
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
with torch.no_grad():
result, temp_feats = model(return_loss=False, temp_feats=temp_feats, rescale=True, **data)
return result, temp_feats
def inference_detector(model, img):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# Use torchvision ops for CPU mode instead
for m in model.modules():
if isinstance(m, (RoIPool, RoIAlign)):
if not m.aligned:
# aligned=False is not implemented on CPU
# set use_torchvision on-the-fly
m.use_torchvision = True
warnings.warn('We set use_torchvision=True in CPU mode.')
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
# forward the model
with torch.no_grad():
result = model(return_loss=False,
return_gfl=True,
rescale=True,
**data)[0]
return result
def matting_inference_file(model,
img,
trimap=None,
mask=None,
image_path="input file directly"):
"""Inference image(s) with the model.
Args:
model (nn.Module): The loaded model.
img (str): Image file path.
trimap (str): Trimap file path.
Returns:
np.ndarray: The predicted alpha matte.
"""
assert trimap is not None or mask is not None
cfg = model.cfg
device = next(model.parameters()).device # model device
# remove alpha from test_pipeline
keys_to_remove = ['alpha', 'ori_alpha']
for key in keys_to_remove:
for pipeline in list(cfg.test_pipeline):
if 'key' in pipeline and key == pipeline['key']:
cfg.test_pipeline.remove(pipeline)
if 'keys' in pipeline and key in pipeline['keys']:
pipeline['keys'].remove(key)
if len(pipeline['keys']) == 0:
cfg.test_pipeline.remove(pipeline)
if 'meta_keys' in pipeline and key in pipeline['meta_keys']:
pipeline['meta_keys'].remove(key)
# build the data pipeline
test_pipeline = cfg.test_pipeline[2:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(merged=img,
mask=mask,
ori_mask=mask,
trimap=trimap,
ori_trimap=trimap,
ori_merged=img.copy(),
merged_path=image_path,
merged_ori_shape=img.shape)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
# forward the model
with torch.no_grad():
result = model(test_mode=True, **data)
return result
def fake_data(model, input=(3, 800, 1333)):
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
if isinstance(input, (list, tuple)):
img = np.random.random(input[::-1])
elif isinstance(input, str):
img = input
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
return data
def inference():
score_cache = deque()
scores_sum = 0
cur_time = time.time()
while True:
cur_windows = []
while len(cur_windows) == 0:
if len(frame_queue) == sample_length:
cur_windows = list(np.array(frame_queue))
if data['img_shape'] is None:
data['img_shape'] = frame_queue.popleft().shape[:2]
cur_data = data.copy()
cur_data['imgs'] = cur_windows
cur_data = test_pipeline(cur_data)
cur_data = collate([cur_data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
cur_data = scatter(cur_data, [device])[0]
with torch.no_grad():
scores = model(return_loss=False, **cur_data)[0]
score_cache.append(scores)
scores_sum += scores
if len(score_cache) == average_size:
scores_avg = scores_sum / average_size
num_selected_labels = min(len(label), 5)
scores_tuples = tuple(zip(label, scores_avg))
scores_sorted = sorted(scores_tuples,
key=itemgetter(1),
reverse=True)
results = scores_sorted[:num_selected_labels]
result_queue.append(results)
scores_sum -= score_cache.popleft()
if inference_fps > 0:
# add a limiter for actual inference fps <= inference_fps
sleep_time = 1 / inference_fps - (time.time() - cur_time)
if sleep_time > 0:
time.sleep(sleep_time)
cur_time = time.time()
camera.release()
cv2.destroyAllWindows()
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
results = [None for _ in range(len(self.dataset))]
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in range(runner.rank, len(self.dataset), runner.world_size):
data = self.dataset[idx]
data_gpu = scatter(collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(return_loss=False,
rescale=True,
**data_gpu)
# result is N polygons [#N, 72d + 1d]
_bbox_lists = []
for j in range(result.shape[0]):
_bbox_lists.append(
dict(points=result[j, :-1].reshape((-1, 2)).tolist(),
score=result[j, -1]))
results[idx] = dict(image_id=self.dataset.img_ids[idx],
bboxes=_bbox_lists)
batch_size = runner.world_size
if runner.rank == 0:
for _ in range(batch_size):
prog_bar.update()
if runner.rank == 0:
dist.barrier()
for i in range(1, runner.world_size):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_reads = mmcv.load(tmp_file)
tmp_results = tmp_reads['results']
for idx in range(i, len(results), runner.world_size):
results[idx] = tmp_results[idx]
os.remove(tmp_file)
self.evaluate(runner, results)
else:
tmp_file = osp.join(runner.work_dir,
'temp_{}.pkl'.format(runner.rank))
writes = {'results': results}
mmcv.dump(writes, tmp_file)
dist.barrier()
dist.barrier()
def inference_detector_batch(model, imgs):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
assert len(imgs) > 1, "batch 大于 1"
import time
from collections import defaultdict
s = time.time()
result = defaultdict()
def parse(f):
idx, data = f.result()
result[idx] = data
def read_data(img, idx):
data = dict(img=img)
data = test_pipeline(data)
return idx, data
batch_size = len(imgs)
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadImages()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
datas = []
executor = ThreadPoolExecutor()
for i, img in enumerate(imgs):
executor.submit(read_data, img, i).add_done_callback(parse)
executor.shutdown()
for i in range(batch_size):
datas.append(result[i])
data = scatter(collate(datas, samples_per_gpu=batch_size), [device])[0]
# print("read data time: ", time.time() - s)
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
return result
def inference_model_batch(model, images_nps, topn=5):
"""Inference image(s) with the classifier.
Args:
model (nn.Module): The loaded classifier.
img (str/ndarray): The image filename or loaded image.
Returns:
result (list of dict): The classification results that contains
`class_name`, `pred_label` and `pred_score`.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile':
cfg.data.test.pipeline.pop(0)
test_pipeline = Compose(cfg.data.test.pipeline)
with torch.no_grad():
inference_results = []
for images_batch in sly.batched(images_nps, g.batch_size):
data = [dict(img=img) for img in images_batch]
data = [test_pipeline(row) for row in data]
data = collate(data, samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
batch_scores = np.asarray(model(return_loss=False, **data))
batch_top_indexes = batch_scores.argsort(axis=1)[:,
-topn:][:, ::-1]
for scores, top_indexes in zip(batch_scores, batch_top_indexes):
inference_results.append({
'label':
top_indexes.astype(int).tolist(),
'score':
scores[top_indexes].astype(float).tolist(),
'class':
np.asarray(model.CLASSES)[top_indexes].tolist()
})
return inference_results
def inference_detector_huge_image(model, img, split_cfg, merge_cfg):
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = [LoadPatch()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# assert model
is_cuda = next(model.parameters()).is_cuda
if not is_cuda:
# Use torchvision ops for CPU mode instead
for m in model.modules():
if isinstance(m, (RoIPool, RoIAlign)):
if not m.aligned:
# aligned=False is not implemented on CPU
# set use_torchvision on-the-fly
m.use_torchvision = True
warnings.warn('We set use_torchvision=True in CPU mode.')
# generate patch windows
img = mmcv.imread(img)
height, width = img.shape[:2]
sizes, steps = parse_split_cfg(split_cfg)
windows = get_windows(width, height, sizes, steps)
# detection loop
results = []
prog_bar = mmcv.ProgressBar(len(windows))
for win in windows:
data = dict(img=img)
data['patch_win'] = win.tolist()
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'][0].data
# forward the model
with torch.no_grad():
results.append(model(return_loss=False, rescale=True, **data))
prog_bar.update()
# merge results
print()
print('Merge patch results!!')
results = merge_patch_results(results, windows, merge_cfg)
return results
def reid_encoder(image, boxes):
image_patches = []
for box in boxes:
patch = extract_image_patch(image, box)
if patch is None:
print("WARNING: Failed to extract image patch: %s." % str(box))
patch = np.random.uniform(0., 255.,
input_size).astype(np.uint8)
data = patch_transform(dict(img=patch))
image_patches.append(data)
if len(image_patches) == 0:
return torch.zeros((0, 512)).to(device)
data = scatter(
collate(image_patches, samples_per_gpu=len(image_patches)),
[device])[0]
with torch.no_grad():
return extractor(**data)
def inference(config_file, checkpoint, classes, args):
cfg = Config.fromfile(config_file)
model = init_model(cfg, checkpoint, device=args.device)
model.CLASSES = classes
# build the data pipeline
if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile':
cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile'))
if cfg.data.test.type in ['CIFAR10', 'CIFAR100']:
# The image shape of CIFAR is (32, 32, 3)
cfg.data.test.pipeline.insert(1, dict(type='Resize', size=32))
data = dict(img_info=dict(filename=args.img), img_prefix=None)
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
resolution = tuple(data['img'].shape[1:])
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [args.device])[0]
# forward the model
result = {'resolution': resolution}
with torch.no_grad():
if args.inference_time:
time_record = []
for _ in range(10):
start = time()
scores = model(return_loss=False, **data)
time_record.append((time() - start) * 1000)
result['time_mean'] = np.mean(time_record[1:-1])
result['time_std'] = np.std(time_record[1:-1])
else:
scores = model(return_loss=False, **data)
pred_score = np.max(scores, axis=1)[0]
pred_label = np.argmax(scores, axis=1)[0]
result['pred_label'] = pred_label
result['pred_score'] = float(pred_score)
result['pred_class'] = model.CLASSES[result['pred_label']]
result['model'] = config_file.stem
return result
def inference_detector(model, img, cfg):
"""Inference image(s) with the detector.
Args:
model (nn.Module): The loaded detector.
imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
images.
Returns:
If imgs is a str, a generator will be returned, otherwise return the
detection results directly.
"""
device = next(model.parameters()).device # model device
# prepare data
if isinstance(img, np.ndarray):
# directly add img
data = dict(img=img)
cfg = cfg.copy()
# set loading pipeline type
cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
else:
# add information into dict
data = dict(img_info=dict(filename=img), img_prefix=None)
# build the data pipeline
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
# just get the actual data from DataContainer
data['img_metas'] = [img_metas.data[0] for img_metas in data['img_metas']]
data['img'] = [img.data[0] for img in data['img']]
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
for m in model.modules():
assert not isinstance(
m, RoIPool
), 'CPU inference with RoIPool is not supported currently.'
# forward the model
with torch.no_grad():
t = time()
result = model(return_loss=False, rescale=True, **data)[0]
t2 = time()-t
return result, t2
def regress_and_classify(self, img, tracklets):
cfg = self.model.cfg
device = next(self.model.parameters()).device
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
data = dict(img=img)
data = test_pipeline(data)
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
proposals = np.array([[
*(tracklet.last_detection.box[:]), tracklet.last_detection.score
] for tracklet in tracklets])
proposals[:, 0:4] = proposals[:, 0:4] * data['img_meta'][0][0][
'scale_factor']
proposals_tensor = torch.tensor(proposals).cuda(device)
with torch.no_grad():
x = self.model.extract_feat(data['img'][0])
# Here I re-implement BBoxTestMixin.simple_test_bboxes().
# It seems that the bboxes shouldn't have been rescaled when calling get_det_bboxes()?
rois = bbox2roi([proposals_tensor])
roi_feats = self.model.bbox_roi_extractor(
x[:len(self.model.bbox_roi_extractor.featmap_strides)], rois)
if self.model.with_shared_head:
roi_feats = self.model.shared_head(roi_feats)
cls_score, bbox_pred = self.model.bbox_head(roi_feats)
bboxes, scores = self.model.bbox_head.get_det_bboxes(
rois,
cls_score,
bbox_pred,
data['img_meta'][0][0]['img_shape'],
data['img_meta'][0][0]['scale_factor'],
# I mean here
rescale=True,
cfg=None)
bboxes = bboxes.view(-1, 81, 4)
bboxes = torch.cat((bboxes[:, 1, :], scores[:, 1:2]), dim=1)
labels = torch.tensor([0] * bboxes.shape[0])
bbox_results = bbox2result(bboxes, labels,
self.model.bbox_head.num_classes)[0]
return bbox_results
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
results = [None for _ in range(len(self.dataset))]
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in range(runner.rank, len(self.dataset), runner.world_size):
data = self.dataset[idx]
data_gpu = scatter(collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
if not isinstance(data_gpu['img'], list):
data_gpu['img'] = [data_gpu['img']]
if not isinstance(data_gpu['img_meta'][0], list):
data_gpu['img_meta'] = [data_gpu['img_meta']]
# compute output
with torch.no_grad():
result = runner.model(return_loss=False,
rescale=True,
**data_gpu)
results[idx] = result
# batch_size = runner.world_size
# if runner.rank == 0:
# for _ in range(batch_size):
# prog_bar.update()
if runner.rank == 0:
print('\n')
dist.barrier()
for i in range(1, runner.world_size):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_results = mmcv.load(tmp_file)
for idx in range(i, len(results), runner.world_size):
results[idx] = tmp_results[idx]
os.remove(tmp_file)
self.evaluate(runner, results)
else:
tmp_file = osp.join(runner.work_dir,
'temp_{}.pkl'.format(runner.rank))
mmcv.dump(results, tmp_file)
dist.barrier()
dist.barrier()
def inference_detector(model, pcd):
"""Inference point cloud with the detector.
Args:
model (nn.Module): The loaded detector.
pcd (str): Point cloud files.
Returns:
tuple: Predicted results and data from pipeline.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = deepcopy(cfg.data.test.pipeline)
test_pipeline = Compose(test_pipeline)
box_type_3d, box_mode_3d = get_box_type(cfg.data.test.box_type_3d)
data = dict(
pts_filename=pcd,
box_type_3d=box_type_3d,
box_mode_3d=box_mode_3d,
# for ScanNet demo we need axis_align_matrix
ann_info=dict(axis_align_matrix=np.eye(4)),
sweeps=[],
# set timestamp = 0
timestamp=[0],
img_fields=[],
bbox3d_fields=[],
pts_mask_fields=[],
pts_seg_fields=[],
bbox_fields=[],
mask_fields=[],
seg_fields=[])
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device.index])[0]
else:
# this is a workaround to avoid the bug of MMDataParallel
data['img_metas'] = data['img_metas'][0].data
data['points'] = data['points'][0].data
# forward the model
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
return result, data
def sample_img2img_model(model, image_path, target_domain=None, **kwargs):
"""Sampling from translation models.
Args:
model (nn.Module): The loaded model.
image_path (str): File path of input image.
style (str): Target style of output image.
Returns:
Tensor: Translated image tensor.
"""
assert isinstance(model, BaseTranslationModel)
# get source domain and target domain
if target_domain is None:
target_domain = model._default_domain
source_domain = model.get_other_domains(target_domain)[0]
cfg = model._cfg
device = next(model.parameters()).device # model device
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
# prepare data
data = dict()
# dirty code to deal with test data pipeline
data['pair_path'] = image_path
data[f'img_{source_domain}_path'] = image_path
data[f'img_{target_domain}_path'] = image_path
data = test_pipeline(data)
if device.type == 'cpu':
data = collate([data], samples_per_gpu=1)
data['meta'] = []
else:
data = scatter(collate([data], samples_per_gpu=1), [device])[0]
source_image = data[f'img_{source_domain}']
# forward the model
with torch.no_grad():
results = model(source_image,
test_mode=True,
target_domain=target_domain,
**kwargs)
output = results['target']
return output
def detect(self, pcd):
#self.data['pts_filename']=pcd
#data = self.datapipelinefromlidarfile(self.model.cfg, pcd)
data = self.datafrompcd(self.model.cfg, pcd)
device = next(self.model.parameters()).device # model device
if next(self.model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device.index])[0]
# forward the model
with torch.no_grad():
result = self.model(return_loss=False, rescale=True, **data)
#return result, data
boxes_3d = result[0]['boxes_3d'].tensor.numpy()
scores_3d = result[0]['scores_3d'].numpy()
labels_3d = result[0]['labels_3d'].numpy()
return {'boxes': boxes_3d, 'scores': scores_3d, 'classes': labels_3d}
def inference():
if len(frame_queue) != sample_length:
# Do no inference when there is no enough frames
return False, None
cur_windows = list(np.array(frame_queue))
img = frame_queue.popleft()
if data['img_shape'] is None:
data['img_shape'] = img.shape[:2]
cur_data = data.copy()
cur_data['imgs'] = cur_windows
cur_data = test_pipeline(cur_data)
cur_data = collate([cur_data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
cur_data = scatter(cur_data, [device])[0]
with torch.no_grad():
scores = model(return_loss=False, **cur_data)[0]
return True, scores
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
# rank = int(os.environ['RANK'])
num_gpus = torch.cuda.device_count()
if runner.rank >= num_gpus:
return
eval_world_size = num_gpus
runner.model.eval()
results = [None for _ in range(len(self.dataset))]
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in range(runner.rank, len(self.dataset), eval_world_size):
data = self.dataset[idx]
data_gpu = scatter(collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(return_loss=False,
rescale=True,
**data_gpu)
results[idx] = result
batch_size = eval_world_size
for _ in range(batch_size):
prog_bar.update()
if runner.rank == 0:
print('\n')
self._barrier(runner.rank, eval_world_size)
for i in range(1, eval_world_size):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_results = mmcv.load(tmp_file)
for idx in range(i, len(results), eval_world_size):
results[idx] = tmp_results[idx]
os.remove(tmp_file)
self.evaluate(runner, results)
else:
tmp_file = osp.join(runner.work_dir,
'temp_{}.pkl'.format(runner.rank))
mmcv.dump(results, tmp_file)
self._barrier(runner.rank, eval_world_size)
self._barrier(runner.rank, eval_world_size)
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
if runner.rank == 0:
outputs = defaultdict(list)
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in range(len(self.dataset)):
data = self.dataset[idx]
data_gpu = scatter(
collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(
return_loss=False, rescale=True, **data_gpu)
outputs['bbox_results'].append(result['bbox_results'])
outputs['new_bbox_results'].append(result['new_bbox_results'])
outputs['track_results'].append(result['track_results'])
if 'segm_results' in result.keys():
outputs['bbox_results'][-1] = (outputs['bbox_results'][-1], result['segm_results'])
outputs['segm_track_results'].append(result['segm_track_results'])
prog_bar.update()
out_name = '{}/tmp.pkl'.format(runner.work_dir)
print('\nwriting results to {}'.format(out_name))
mmcv.dump(outputs, out_name)
# bbox
result_files = results2json(self.dataset, outputs['bbox_results'],
out_name)
coco_eval(result_files, ['bbox'], self.val_cfg.ann_file)
# Box tracking
result_files = results2json(self.dataset,
outputs['new_bbox_results'], out_name)
coco_eval(result_files, ['bbox'], self.val_cfg.ann_file)
print("Evaluating box tracking...")
mdat_eval_out = mdat_eval(outputs['track_results'], self.dataset, out_name, ann_file=self.val_cfg.ann_file)
if 'segm_track_results' in outputs.keys():
print("Evaluating seg tracking...")
out = mdat_eval(outputs['segm_track_results'], self.dataset, out_name, ann_file=self.val_cfg.ann_file, with_mask=True)
out = {'seg_' + k: v for k, v in out.items()}
mdat_eval_out.update(out)
runner.log_buffer.output.update(mdat_eval_out)
runner.log_buffer.ready = True
def _inference(self, filename):
ss = time.time()
data = dict(img=filename)
data = self.test_pipeline(data)
data = collate([data], samples_per_gpu=1)
data = scatter(data, [self.device])[0]
with torch.no_grad():
pred = self.model(return_loss=False, rescale=True, **data)
#print(len(pred[0][0]), len(pred[0][1]), 'haha')
#print(pred[0][0][0].shape, pred[0][0][1].shape)
bbox = pred[0][0][1]
area = (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1])
valid = area > 100
bbox = bbox[valid]
pred = pred[0][1]
tot = len(pred[1])
res = pred[1][0].astype(np.int32)
for i in range(1, tot):
res = res + pred[1][i].astype(np.int32)
tot_bone = len(pred[0])
res_bone = res * 0
for i in range(0, tot_bone):
res_bone = res_bone + pred[0][i].astype(np.int32)
res = res > 0
res = res.astype(np.uint8)
res_bone = res_bone > 0
res_bone = res_bone.astype(np.uint8)
res = res * res_bone
res = res * 255
return res, bbox
def inference_model(model, img, topn=5):
"""Inference image(s) with the classifier.
Args:
model (nn.Module): The loaded classifier.
img (str/ndarray): The image filename or loaded image.
Returns:
result (list of dict): The classification results that contains
`class_name`, `pred_label` and `pred_score`.
"""
cfg = model.cfg
device = next(model.parameters()).device # model device
# build the data pipeline
if isinstance(img, str):
if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile':
cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile'))
data = dict(img_info=dict(filename=img), img_prefix=None)
else:
if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile':
cfg.data.test.pipeline.pop(0)
data = dict(img=img)
test_pipeline = Compose(cfg.data.test.pipeline)
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
# forward the model
with torch.no_grad():
scores = model(return_loss=False, **data)
model_out = scores[0]
top_scores = model_out[model_out.argsort()[-topn:]][::-1]
top_labels = model_out.argsort()[-topn:][::-1]
result = []
for label, score in zip(top_labels, top_scores):
result.append({
'label': int(label),
'score': float(score),
'class': model.CLASSES[label]
})
return result
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.model.eval()
runner.model.module.data = self.cfg
runner.model.module.dataset = self.dataset
runner.model.module.out = '{}/output'.format(self.work_dir)
results = [None for _ in range(len(self.dataset))]
dist_indices = self.get_dist_indices(runner.world_size)
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(self.dataset))
for idx in dist_indices[runner.rank]:
data = self.dataset[idx]
data_gpu = scatter(
collate([data], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
# compute output
with torch.no_grad():
result = runner.model(
return_loss=False, rescale=True, **data_gpu)
results[idx] = result
if runner.rank == 0:
for _ in range(runner.world_size):
prog_bar.update()
if runner.rank == 0:
print('\n')
dist.barrier()
for i in range(1, runner.world_size):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_results = mmcv.load(tmp_file)
for idx in dist_indices[i]:
results[idx] = tmp_results[idx]
os.remove(tmp_file)
self.evaluate(runner, results)
else:
tmp_file = osp.join(runner.work_dir,
'temp_{}.pkl'.format(runner.rank))
mmcv.dump(results, tmp_file)
dist.barrier()
dist.barrier()
