def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
"""Test model with multiple gpus.
This method tests model with multiple gpus and collects the results
under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
it encodes results to gpu tensors and use gpu communication for results
collection. On cpu mode it saves the results on different gpus to 'tmpdir'
and collects them by the rank 0 worker.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): Pytorch data loader.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode.
gpu_collect (bool): Option to use either gpu or cpu to collect results.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
time.sleep(2) # This line can prevent deadlock problem in some cases.
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
# encode mask results
if isinstance(result, tuple) and len(result) == 2:
bbox_results, mask_results = result
encoded_mask_results = encode_mask_results(mask_results)
result = bbox_results, encoded_mask_results
elif isinstance(result, tuple) and len(result) == 3:
bbox_results, mask_results, offset_results = result
if mask_results is not None:
encoded_mask_results = encode_mask_results(mask_results)
result = bbox_results, encoded_mask_results, offset_results
else:
# only pred offset
result = bbox_results, offset_results
elif isinstance(result, tuple) and len(result) == 4:
bbox_results, mask_results, offset_results, height_results = result
encoded_mask_results = encode_mask_results(mask_results)
result = bbox_results, encoded_mask_results, offset_results, height_results
results.append(result)
if rank == 0:
batch_size = (len(data['img_meta'].data) if 'img_meta' in data else
len(data['img_metas'][0].data))
for _ in range(batch_size * world_size):
prog_bar.update()
# collect results from all ranks
if gpu_collect:
results = collect_results_gpu(results, len(dataset))
else:
results = collect_results_cpu(results, len(dataset), tmpdir)
return results
def evaluate_full_dataset(
self, data_loader: torch.utils.data.DataLoader) -> Dict[str, Any]:
if self.data_config["backend"] == "fake":
return {"bbox_mAP": 0}
# Will need custom reducer to do this across gpus
prog_bar = ProgressBar(len(data_loader.dataset))
results = []
for i, batch in enumerate(data_loader):
# TODO: modify this to use cpu_only field of DataContainers.
batch["img"] = [self.context.to_device(batch["img"])]
batch = {key: batch[key][0] for key in batch}
with torch.no_grad():
result = self.model(return_loss=False, rescale=True, **batch)
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
batch_size = len(result)
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
eval_kwargs = self.cfg.evaluation
for key in ["interval", "tmpdir", "start", "gpu_collect"]:
eval_kwargs.pop(key, None)
metrics = data_loader.dataset.evaluate(results, **eval_kwargs)
return metrics
def single_gpu_test_analysis(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3,
analysis_scale=1.0):
model.eval()
results = []
dataset = data_loader.dataset
# print("dataset type", type(dataset))
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
if i >= 10:
break
with torch.no_grad():
# print("data keys", data.keys())
# print("data img_metas", data['img_metas'])
result = model(return_loss=True,
rescale=True,
**data,
batch_idx=i,
analysis_scale=analysis_scale)
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def encode_segms(results, dataset):
new_results = []
for result in results:
# encode mask results
if isinstance(result, tuple):
if len(result) == 2:
bbox_results, mask_results = result
elif len(result) == 3:
bbox_results, mask_results, stuff_results = result
else:
raise ValueError('Wrong result.')
encoded_mask_results = encode_mask_results(mask_results)
if len(result) == 2:
result = bbox_results, encoded_mask_results
elif len(result) == 3:
stuff_segms = [[] for _ in range(len(dataset.seg_ids))]
stuff_map = stuff_results['stuff_map']
stuff_img_shape = stuff_results['img_shape']
stuff_map = mmcv.imresize(stuff_map,
stuff_img_shape,
interpolation='nearest')
unique_stuffs = np.unique(stuff_map)
for j in unique_stuffs:
stuff_class_mask = (stuff_map == j).astype(np.uint8)
rle = mask_util.encode(
np.array(stuff_class_mask[:, :, np.newaxis],
order='F'))[0]
stuff_segms[j].append(rle)
result = bbox_results, encoded_mask_results, stuff_segms
else:
raise ValueError('Wrong result.')
new_results.append(result)
return new_results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
print('clw: using single_gpu_test() !!')
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
########### clw note: for debug
# for idx, item in enumerate(result[0]):
# for row in item:
# print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
# if row[2] - row[0] == 0 or row[3] - row[1] == 0:
# print('aaaa')
#########
#batch_size = len(result)
batch_size = data["img"][0].size(0) # clw modify: TSD model wrong
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def multi_gpu_test(model, data_loader, tmpdir=None):
model.eval()
results = []
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
if rank == 0:
batch_size = len(result)
for _ in range(batch_size * world_size):
prog_bar.update()
# collect results from all ranks
results = collect_results(results, len(dataset), tmpdir)
return results
def postprocess(result, img_meta, num_classes=80, rescale=True):
det_bboxes = result['boxes']
det_labels = result['labels']
det_masks = result.get('masks', None)
if rescale:
img_h, img_w = img_meta[0]['ori_shape'][:2]
scale = img_meta[0]['scale_factor']
det_bboxes[:, :4] /= scale
else:
img_h, img_w = img_meta[0]['img_shape'][:2]
det_bboxes[:, 0:4:2] = np.clip(det_bboxes[:, 0:4:2], 0, img_w - 1)
det_bboxes[:, 1:4:2] = np.clip(det_bboxes[:, 1:4:2], 0, img_h - 1)
bbox_results = bbox2result(det_bboxes, det_labels, num_classes)
if det_masks is not None:
segm_results = mask2result(
det_bboxes,
det_labels,
det_masks,
num_classes,
mask_thr_binary=0.5,
img_size=(img_h, img_w))
segm_results = encode_mask_results(segm_results)
return bbox_results, segm_results
return bbox_results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
flops_ori = AverageMeter()
flops_real = AverageMeter()
for i, data in enumerate(data_loader):
with torch.no_grad():
result, real_f, ori_f = model(return_loss=False,
rescale=True,
**data)
batch_size = len(result)
flops_ori.update(ori_f.item(), batch_size)
flops_real.update(real_f.item(), batch_size)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
# show the flops
print('\nAverage FLOPs: real = {}M, original = {}M'.format(
flops_real.avg / 1024 / 1024, flops_ori.avg / 1024 / 1024))
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
filename = 'result.txt'
with open(filename, 'w') as f:
f.write(result)
# print('in api test result: ')
# print(type(result))
# print(result)
batch_size = len(result)
if show or out_dir:
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 = osp.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)
with open(filename, 'w', 'UTF-8') as f:
f.write(result)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3,
save_crop_image = True):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
if show or out_dir:
img_tensor = data['img'][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 img, img_meta in 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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
model.module.show_result(
img_show,
result,
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result, tuple):
bbox_results, mask_results = result
if save_crop_image:
crop_image(result, INPUT_PATH, OUTPUT_PATH,i)
encoded_mask_results = encode_mask_results(mask_results)
result = bbox_results, encoded_mask_results
results.append(result)
batch_size = len(data['img_metas'][0].data)
for _ in range(batch_size):
prog_bar.update()
return results
def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
"""Test model with multiple gpus.
This method tests model with multiple gpus and collects the results
under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
it encodes results to gpu tensors and use gpu communication for results
collection. On cpu mode it saves the results on different gpus to 'tmpdir'
and collects them by the rank 0 worker. 'gpu_collect=True' is not
supported for now.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): Pytorch data loader.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode. Defaults to None.
gpu_collect (bool): Option to use either gpu or cpu to collect results.
Defaults to False.
Returns:
dict[str, list]: The prediction results.
"""
model.eval()
results = defaultdict(list)
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
time.sleep(2) # This line can prevent deadlock problem in some cases.
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
for key in result:
if 'mask' in key:
result[key] = encode_mask_results(result[key])
for k, v in result.items():
results[k].append(v)
if rank == 0:
batch_size = data['img'][0].size(0)
for _ in range(batch_size * world_size):
prog_bar.update()
# collect results from all ranks
if gpu_collect:
raise NotImplementedError
else:
results = collect_results_cpu(results, tmpdir)
return results
def single_gpu_test(model,
img_data_list,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
for (img, img_meta) in img_data_list:
data = {
"img_metas": [DC([[img_meta]], cpu_only=True)],
"img": [img.unsqueeze(0).contiguous()]
}
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
if show or out_dir:
img_tensor = data['img'][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 img, img_meta in 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 = osp.join(out_dir,
osp.basename(img_meta['ori_filename']))
else:
out_file = None
show_result(img_show,
result,
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result, tuple):
bbox_results, mask_results = result
encoded_mask_results = encode_mask_results(mask_results)
result = bbox_results, encoded_mask_results
results.append(result)
return results
def test_instance_segmentation_evaluation():
pred_bbox = [
np.array([[11, 10, 20, 50, 0.8], [31, 10, 40, 50, 0.8]]),
np.array([[51, 10, 60, 15, 0.7]])
]
person1_mask = np.zeros((60, 80), dtype=bool)
person1_mask[20:50, 11:20] = True
person2_mask = np.zeros((60, 80), dtype=bool)
person2_mask[20:50, 31:40] = True
dog_mask = np.zeros((60, 80), dtype=bool)
dog_mask[10:15, 51:60] = True
pred_mask = [[person1_mask, person2_mask], [
dog_mask,
]]
results = [{'ins_results': (pred_bbox, encode_mask_results(pred_mask))}]
tmp_dir = tempfile.TemporaryDirectory()
pan_ann_file = osp.join(tmp_dir.name, 'panoptic.json')
ins_ann_file = osp.join(tmp_dir.name, 'instance.json')
_create_panoptic_gt_annotations(pan_ann_file)
_create_instance_segmentation_gt_annotations(ins_ann_file)
dataset = CocoPanopticDataset(ann_file=pan_ann_file,
ins_ann_file=ins_ann_file,
seg_prefix=tmp_dir.name,
pipeline=[])
dataset.THING_CLASSES = ['person', 'dog']
dataset.STUFF_CLASSES = ['wall']
dataset.CLASSES = dataset.THING_CLASSES + dataset.STUFF_CLASSES
parsed_results = dataset.evaluate(results, metric=['segm', 'bbox'])
# Here is the results for instance segmentation:
# {
# 'segm_mAP': 0.5, 'segm_mAP_50': 0.626, 'segm_mAP_75': 0.5,
# 'segm_mAP_s': 0.5, 'segm_mAP_m': -1.0, 'segm_mAP_l': -1.0,
# 'segm_mAP_copypaste': '0.500 0.626 0.500 0.500 -1.000 -1.000',
# 'bbox_mAP': 0.564, 'bbox_mAP_50': 0.626, 'bbox_mAP_75': 0.626,
# 'bbox_mAP_s': 0.564, 'bbox_mAP_m': -1.0, 'bbox_mAP_l': -1.0,
# 'bbox_mAP_copypaste': '0.564 0.626 0.626 0.564 -1.000 -1.000'
# }
assert np.isclose(parsed_results['segm_mAP'], 0.5)
assert np.isclose(parsed_results['bbox_mAP'], 0.564)
def single_gpu_test(model, data_loader, show=False):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=not show, **data)
if show:
model.module.show_result(data, result, dataset.img_norm_cfg)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
batch_size = len(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
eval_results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
if show or out_dir:
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, :]
if img_show.shape[-1] > 3:
img_show = img[:h, :w, :3][:, :, ::-1]
else:
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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
model.module.show_result(img_show,
result[i][:2],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result[0], tuple):
if len(result[0]) == 3:
result = [(bbox_results, encode_mask_results(mask_results),
poly_points)
for bbox_results, mask_results, poly_points in result
]
else:
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
eval_results.extend([res[:2] for res in result])
for _ in range(batch_size):
prog_bar.update()
return results, eval_results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
is_kie=False,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
if show or out_dir:
if is_kie:
img_tensor = data['img'].data[0]
if img_tensor.shape[0] != 1:
raise KeyError('Visualizing KIE outputs in batches is'
'currently not supported.')
gt_bboxes = data['gt_bboxes'].data[0]
img_metas = data['img_metas'].data[0]
imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
for i, img in enumerate(imgs):
h, w, _ = img_metas[i]['img_shape']
img_show = img[:h, :w, :]
if out_dir:
out_file = osp.join(out_dir,
img_metas[i]['ori_filename'])
else:
out_file = None
model.module.show_result(img_show,
result[i],
gt_bboxes[i],
show=show,
out_file=out_file)
else:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test_rotate_rect_img(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
print('clw: using single_gpu_test_rotate_rect_img() !!')
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
########### clw note: for debug
# for idx, item in enumerate(result[0]):
# if item.size == 0:
# print('111')
# for row in item:
# print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
# if row[2] - row[0] == 0 or row[3] - row[1] == 0:
# print('aaaa')
#########
##
img_name = data['img_metas'][0].data[0][0]['ori_filename']
# origin_name = img_name.split('CAM')[0] + 'CAM' + img_name.split('CAM')[1][0] + '.jpg'
# data['img_metas'][0].data[0][0]['ori_filename'] = origin_name
# data['img_metas'][0].data[0][0]['filename'] = data['img_metas'][0].data[0][0]['filename'].rsplit('/', 1)[0] + '/' + origin_name
aaa = img_name[:-4].split('_')[-9:]
bbb = [float(a) for a in aaa]
M_perspective_inv = np.array(bbb).reshape(3, 3)
for i in range(len(result[0])):
ddd = []
ccc = result[0][i][:, :4] # (n, 4)
if ccc.size == 0:
continue
for xyxy in ccc:
x1 = xyxy[0]
y1 = xyxy[1]
x2 = xyxy[2]
y2 = xyxy[3]
cnt = np.array(((x1, y1), (x1, y2), (x2, y2), (x2, y1)))
ddd.append(cnt)
ddd = np.array(ddd)
#
fff = []
src_pts = cv2.perspectiveTransform(ddd, M_perspective_inv)
for cnt in src_pts:
rect = cv2.boundingRect(cnt)
x1 = rect[0]
y1 = rect[1]
x2 = rect[0] + rect[2]
y2 = rect[1] + rect[3]
ggg = np.array((x1, y1, x2, y2))
fff.append(ggg)
fff = np.array(fff)
result[0][
i][:, :
4] = fff # result[0][i] = np.concatenate((fff, result[0][i][:, 4]), axis=1)
##
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
is_kie=False,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for data in data_loader:
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
if show or out_dir:
if is_kie:
img_tensor = data['img'].data[0]
if img_tensor.shape[0] != 1:
raise KeyError('Visualizing KIE outputs in batches is'
'currently not supported.')
gt_bboxes = data['gt_bboxes'].data[0]
img_metas = data['img_metas'].data[0]
must_keys = ['img_norm_cfg', 'ori_filename', 'img_shape']
for key in must_keys:
if key not in img_metas[0]:
raise KeyError(
f'Please add {key} to the "meta_keys" in config.')
# for no visual model
if np.prod(img_tensor.shape) == 0:
imgs = []
for img_meta in img_metas:
try:
img = mmcv.imread(img_meta['filename'])
except Exception as e:
print(f'Load image with error: {e}, '
'use empty image instead.')
img = np.ones(img_meta['img_shape'],
dtype=np.uint8)
imgs.append(img)
else:
imgs = tensor2imgs(img_tensor,
**img_metas[0]['img_norm_cfg'])
for i, img in enumerate(imgs):
h, w, _ = img_metas[i]['img_shape']
img_show = img[:h, :w, :]
if out_dir:
out_file = osp.join(out_dir,
img_metas[i]['ori_filename'])
else:
out_file = None
model.module.show_result(img_show,
result[i],
gt_bboxes[i],
show=show,
out_file=out_file)
else:
img_tensor, img_metas, img_norm_cfg = \
retrieve_img_tensor_and_meta(data)
if img_tensor.size(1) == 1:
imgs = tensor2grayimgs(img_tensor, **img_norm_cfg)
else:
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
for j, (img, img_meta) in enumerate(zip(imgs, img_metas)):
img_shape, ori_shape = img_meta['img_shape'], img_meta[
'ori_shape']
img_show = img[:img_shape[0], :img_shape[1]]
img_show = mmcv.imresize(img_show,
(ori_shape[1], ori_shape[0]))
if out_dir:
out_file = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
model.module.show_result(img_show,
result[j],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test_processed_rect_img(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
print('clw: using single_gpu_test_processed_rect_img() !!')
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
########### clw note: for debug
# for idx, item in enumerate(result[0]):
# if item.size == 0:
# print('111')
# for row in item:
# print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
# if row[2] - row[0] == 0 or row[3] - row[1] == 0:
# print('aaaa')
#########
##
img_name = data['img_metas'][0].data[0][0]['ori_filename']
aaa = img_name[:-4].split('_')[-2:]
x_rect_left = int(aaa[0])
y_rect_up = int(aaa[1])
for i in range(len(result[0])):
ddd = []
ccc = result[0][i][:, :4] # (n, 4)
if ccc.size == 0:
continue
for xyxy in ccc:
x1 = xyxy[0] + x_rect_left
y1 = xyxy[1] + y_rect_up
x2 = xyxy[2] + x_rect_left
y2 = xyxy[3] + y_rect_up
cnt = np.array((x1, y1, x2, y2))
ddd.append(cnt)
ddd = np.array(ddd)
result[0][
i][:, :
4] = ddd # result[0][i] = np.concatenate((fff, result[0][i][:, 4]), axis=1)
##
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
batch_size = 1
if show or out_dir:
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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
point_box = model.module.show_result(img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# save predict results into txt
import os
import numpy as np
if model.module.test_cfg.text_dataset_type == 'ICDAR2015':
polys = np.array(point_box).reshape(-1, 8)
image_name = img_metas[0]['ori_filename'][0:-4]
with open(
'{}'.format(
os.path.join(
'evaluate/icdar2015_evalu/res_mask',
'res_{}.txt'.format(image_name))),
'w') as f:
import cv2
for id in range(polys.shape[0]):
# if len(polys[id]) < 8:
# continue
# img = mmcv.imread(img_show)
# img = img.copy()
# cv2.line(img, (polys[id][0], polys[id][1]), (polys[id][2], polys[id][3]), (0, 255, 0), 1, 1)
# cv2.line(img, (polys[id][2], polys[id][3]), (polys[id][4], polys[id][5]), (0, 255, 0), 1, 1)
# cv2.line(img, (polys[id][4], polys[id][5]), (polys[id][6], polys[id][7]), (0, 255, 0), 1, 1)
# cv2.line(img, (polys[id][6], polys[id][7]), (polys[id][0], polys[id][1]), (0, 255, 0), 1, 1)
f.write('{}, {}, {}, {}, {}, {}, {}, {}\n'.format(
int(polys[id][0]), int(polys[id][1]),
int(polys[id][2]), int(polys[id][3]),
int(polys[id][4]), int(polys[id][5]),
int(polys[id][6]), int(polys[id][7])))
# cv2.imwrite(out_file, img)
elif model.module.test_cfg.text_dataset_type == 'CTW1500':
image_name = img_metas[0]['ori_filename'][0:-4]
with open(
'{}'.format(
os.path.join('evaluate/ctw1500_evalu/predict1',
'{}.txt'.format(image_name))),
'w') as f:
for id in range(len(point_box)):
if len(point_box[id]) < 6:
continue
values = [int(v) for v in point_box[id]]
for ii in range(len(values)):
f.write('{}'.format(values[ii]))
if ii < len(values) - 1:
f.write(', ')
f.write('\n')
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
#save result to json file as gt (auto label)
# img_file = data["img_metas"][0].data[0][0]["filename"]
# json_file = img_file.replace(".jpg", ".json")
# to_dump_dict = {}
# to_dump_dict["flags"] = {}
# to_dump_dict["imagePath"] = data["img_metas"][0].data[0][0]["filename"].split("/")[-1]
# to_dump_dict["imageHeight"] = data["img_metas"][0].data[0][0]["ori_shape"][0]
# to_dump_dict["imageWidth"] = data["img_metas"][0].data[0][0]["ori_shape"][1]
# to_dump_dict["version"] = "3.11.2"
# to_dump_dict["imageData"] = None
# to_dump_dict["lineColor"] = [0, 255, 0, 128]
# to_dump_dict["fillColor"] = [255, 0, 0, 128]
# to_dump_dict["shapes"] = []
# classmap = {0:"plate", 1:"headstock", 2:"tailstock", 3:"car", 4:"side_window", 5:"window", 6:"roof", 7:"person", 8:"cycle"}
# for j in range(len(result[0])):
# r = result[0][j]
# obj_num = r.shape[0]
# for jj in range(obj_num):
# x_min = int(r[jj][0])
# y_min = int(r[jj][1])
# x_max = int(r[jj][2])
# y_max = int(r[jj][3])
# score = float(r[jj][4])
# if score < 0.85:
# continue
# #crop car and save:
# #--------------
# import os
# import cv2
# import numpy as np
# if j == 3:
# img_file = data["img_metas"][0].data[0][0]["filename"]
# iimg = cv2.imread(img_file)
# crop_img = mmcv.imcrop(iimg, np.array([x_min,y_min,x_max,y_max]))
# cv2.imwrite(os.path.join("/data/taofuyu/tao_dataset/posture/train_rf/", img_file.split("/")[-1].replace(".jpg", "_"+str(jj)+".jpg")),crop_img)
# #--------------
# points= []
# points.append([x_min, y_min])
# points.append([x_max, y_max])
# to_dump_dict["shapes"].append({"fill_color":None, "line_color":None, "shape_type":"rectangle", "label":classmap[j], "points":points})
# import json
# json.dump(to_dump_dict, open(json_file, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)
batch_size = len(result)
if show or out_dir:
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 = osp.join(out_dir, img_meta['filename'].split("/")[-1])
else:
out_file = None
model.module.show_result(
img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test_processed_rect_crop_img(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
print('clw: using single_gpu_test_processed_rect_crop_img() !!')
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
print(data['img_metas'][0].data[0][0]['ori_filename'])
img_h = data['img'][0].shape[2]
img_w = data['img'][0].shape[3]
with torch.no_grad():
# 否则切图, 4 nums
##############################
overlap_h = 256
overlap_w = 256
crop_h = 2048
crop_w = 2048
step_h = crop_h - overlap_h
step_w = crop_w - overlap_w
nms_iou_thr = model.module.test_cfg['rcnn']['nms']['iou_threshold']
results_crop = [[] for _ in range(len(model.module.CLASSES))]
data['img_metas'][0].data[0][0]['ori_shape'] = (crop_h, crop_w)
data['img_metas'][0].data[0][0]['img_shape'] = (crop_h, crop_w)
data['img_metas'][0].data[0][0]['pad_shape'] = (crop_h, crop_w)
img_tensor_orig = data['img'][0].clone()
for start_h in range(0, img_h - crop_h + 1,
step_h): # imgsz is crop step here,
if start_h + crop_h > img_h: # 如果最后剩下的不到imgsz,则step少一些,保证切的图尺寸不变
start_h = img_h - crop_h
for start_w in range(0, img_w - crop_w + 1, step_w):
if start_w + crop_w > img_w: # 如果最后剩下的不到imgsz,则step少一些,保证切的图尺寸不变
start_w = img_w - crop_w
# crop
print(start_h, start_w)
data['img'][0] = img_tensor_orig[:, :,
start_h:start_h + crop_h,
start_w:start_w + crop_w]
result = model(
return_loss=False, rescale=True, **data
) # result[0]: model.module.CLASSES 个list,每个里面装着(n, 5) ndarray
#result = model(return_loss=False, rescale=False, **data) # clw modify
for idx, item in enumerate(result[0]):
for row in item:
#print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
if row[2] - row[0] == 0 or row[3] - row[1] == 0:
print(
'==================================================================='
)
continue
row[[0, 2]] += start_w
row[[1, 3]] += start_h
results_crop[idx].append(row)
results_afternms = []
for idx, res in enumerate(results_crop):
if len(res) == 0:
results_afternms.append(np.array(
[])) # clw note: it's really important!!
continue
else:
prediction = torch.tensor(res)
boxes, scores = prediction[:, :
4], prediction[:,
4] # boxes (offset by class), scores
i = torchvision.ops.boxes.nms(boxes, scores, nms_iou_thr)
results_afternms.append(prediction[i].numpy())
result = [results_afternms]
##############################
########### clw note: for debug
# for idx, item in enumerate(result[0]):
# if item.size == 0:
# print('111')
# for row in item:
# print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
# if row[2] - row[0] == 0 or row[3] - row[1] == 0:
# print('aaaa')
#########
##
img_name = data['img_metas'][0].data[0][0]['ori_filename']
aaa = img_name[:-4].split('_')[-2:]
x_rect_left = int(aaa[0])
y_rect_up = int(aaa[1])
for i in range(len(result[0])):
ddd = []
if result[0][i].size == 0:
continue
ccc = result[0][i][:, :4] # (n, 4)
for xyxy in ccc:
x1 = xyxy[0] + x_rect_left
y1 = xyxy[1] + y_rect_up
x2 = xyxy[2] + x_rect_left
y2 = xyxy[3] + y_rect_up
cnt = np.array((x1, y1, x2, y2))
ddd.append(cnt)
ddd = np.array(ddd)
result[0][
i][:, :
4] = ddd # result[0][i] = np.concatenate((fff, result[0][i][:, 4]), axis=1)
##
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3,
out=False):
model.eval()
results = []
poly_results = []
filenames = []
sides_list = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result) # 1
if show or out_dir:
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]
height = img_metas[0]['ori_shape'][0]
width = img_metas[0]['ori_shape'][1]
imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
assert len(imgs) == len(img_metas) #1
for j, (img, img_meta) in enumerate(zip(imgs, img_metas)):
h, w, _ = img_meta['img_shape']
filename = img_meta['ori_filename']
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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
filenames.append(filename)
model.module.show_result(img_show,
result[j],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result[0], tuple):
# for i in range(100):
# image = np.array(result[0][1])[:,i,:,:].squeeze()*255.0 # class, proposed, width, height
# im = Image.fromarray(image)
# im = im.convert("L")
# im.save(f"./images_for_check/{i}.png") # why only 100??
# get poly from mask
if batch_size == 1 and out:
poly_results_per_image = []
mask = np.array(result[0][1]) # (1, 105, 852, 852)
bbox = np.array(result[0][0]) # (1, 105, 5)
# print(np.array(mask[0]).shape) # 15, 852, 852)
# print(np.array(bbox[0]).shape) # 15,5
if BINARY:
proposed_num = mask.shape[1]
assert len(
bbox[0]
) == proposed_num, "the num of bbox should equal to that of segm"
#get poly
for k in range(proposed_num):
mask_array = mask[:, k, :, :].squeeze() * 1 #852,852
polygons = Mask(mask_array).polygons()
coords = (polygons.points)
# one proposal should have one blob
if len(coords) == 1:
coords = coords[0]
else:
main_coord_idx = 0
len_main_coord = len(coords[main_coord_idx])
# print("main", len_main_coord)
for l in range(1, len(coords)):
# print(len(coords[l]))
if len(coords[l]) > len_main_coord:
main_coord_idx = l
len_main_coord = len(coords[l])
coords = coords[main_coord_idx]
## get only 4 points
x = []
y = []
for coord in coords:
x.append(coord[0])
y.append(coord[1])
x = np.array(x)
y = np.array(y)
x_max_idx = np.argmax(x)
x_min_idx = np.argmin(x)
y_max_idx = np.argmax(y)
y_min_idx = np.argmin(y)
coords = np.array([[x[x_min_idx], y[x_min_idx]],
[x[y_min_idx], y[y_min_idx]],
[x[x_max_idx], y[x_max_idx]],
[x[y_max_idx], y[y_max_idx]]])
#calculate sides
s = [0] * len(coords)
for m in range(len(coords)):
s[m] = norm(coords[m + 1] -
coords[m]) if m + 1 != len(
coords) else norm(coords[0] -
coords[m])
#averaging max_sides
max_idx = np.argmax(s)
max_length = s[max_idx]
sub_max_idx = (max_idx + 2) % len(coords)
sub_max_length = s[sub_max_idx]
avg_max_length = (max_length + sub_max_length) / 2
if max_idx == 1 or max_idx == 3:
coords = list(coords[-1:]) + list(coords[:-1])
coords = np.array(coords)
if not (s[sub_max_idx] > s[(sub_max_idx + 1) % 4]
and s[sub_max_idx] > s[(sub_max_idx - 1) % 4]):
if (x[x_max_idx] - x[x_min_idx]) > (y[y_max_idx] -
y[y_min_idx]):
avg_max_length = (x[x_max_idx] - x[x_min_idx])
coords = np.array(
[[x[x_min_idx], y[y_min_idx]],
[x[x_max_idx], y[y_min_idx]],
[x[x_max_idx], y[y_max_idx]],
[x[x_min_idx], y[y_max_idx]]])
else:
coords = np.array(
[[x[x_max_idx], y[y_min_idx]],
[x[x_max_idx], y[y_max_idx]],
[x[x_min_idx], y[y_max_idx]],
[x[x_min_idx], y[y_min_idx]]])
avg_max_length = (y[y_max_idx] - y[y_min_idx])
# 회전하면 되는데 귀찮다
image = mask_array * 255.0 # class, proposed, width, height
im = Image.fromarray(image)
im = im.convert("L")
im.save(f"./images_for_check/special{i}_{k}.png"
) # why only 100??
#class check!
if height != width:
raise ValueError(
"image should have the form of square!")
normed_ml = avg_max_length / height
# 35.5, 44, 52, 60,68,76.5
if normed_ml < 35.5 / 852:
class_pred = 1
elif normed_ml < 44 / 852:
class_pred = 2
elif normed_ml < 52 / 852:
class_pred = 3
elif normed_ml < 60 / 852:
class_pred = 4
elif normed_ml < 68 / 852:
class_pred = 5
elif normed_ml < 76.5 / 852:
class_pred = 6
elif normed_ml >= 76.5 / 852:
class_pred = 7
sides_list.append(normed_ml)
# print(normed_ml, class_pred)
poly_results_per_image.append(
(class_pred, bbox[0][k][-1], coords))
else:
num_of_classes = len(mask)
for k in range(num_of_classes):
class_mask = np.array(mask[k]) # 15, 852, 852
class_bbox = np.array(bbox[k]) # 15,5
proposed_num = class_mask.shape[0]
# print("rop", proposed_num)
for l in range(proposed_num):
mask_array = class_mask[l, :, :]
# polygons = Mask(mask_array).polygons()
# coords = polygons.points
# # one proposal should have one blob
# if len(coords) ==1:
# coords = coords[0]
# else:
# main_coord_idx = 0
# len_main_coord = len(coords[main_coord_idx])
# # print("main", len_main_coord)
# for m in range(1,len(coords)):
# # print(len(coords[l]))
# if len(coords[m]) > len_main_coord:
# main_coord_idx = m
# len_main_coord = len(coords[m])
# coords = coords[main_coord_idx]
coords = np.argwhere(mask_array.T == True)
coords = minimum_bounding_rectangle(coords)
# ## get only 4 points
# x = []
# y = []
# for coord in coords:
# x.append(coord[0])
# y.append(coord[1])
# x = np.array(x)
# y = np.array(y)
# x_max_idx = np.argmax(x)
# x_min_idx = np.argmin(x)
# y_max_idx = np.argmax(y)
# y_min_idx = np.argmin(y)
# coords = np.array([[x[x_min_idx], y[x_min_idx]],
# [x[y_min_idx], y[y_min_idx]],
# [x[x_max_idx], y[x_max_idx]],
# [x[y_max_idx], y[y_max_idx]]
# ])
# #calculate sides
# s = [0]*len(coords)
# for n in range(len(coords)):
# s[n] = norm(coords[n+1] - coords[n]) if n+1 != len(coords) else norm(coords[0] - coords[n])
# #averaging max_sides
# max_idx = np.argmax(s)
# max_length = s[max_idx]
# sub_max_idx = (max_idx + 2)%len(coords)
# sub_max_length = s[sub_max_idx]
# avg_max_length = (max_length + sub_max_length)/2
# if max_idx == 1 or max_idx == 3:
# coords = list(coords[-1:]) + list(coords[:-1])
# coords = np.array(coords)
# if not (s[sub_max_idx] > s[(sub_max_idx+1)%4] and s[sub_max_idx]> s[(sub_max_idx-1)%4]):
# if (x[x_max_idx]- x[x_min_idx]) > (y[y_max_idx]- y[y_min_idx]):
# avg_max_length = (x[x_max_idx]- x[x_min_idx])
# coords = np.array([[x[x_min_idx], y[y_min_idx]],
# [x[x_max_idx], y[y_min_idx]],
# [x[x_max_idx], y[y_max_idx]],
# [x[x_min_idx], y[y_max_idx]]
# ])
# else:
# coords = np.array([[x[x_max_idx], y[y_min_idx]],
# [x[x_max_idx], y[y_max_idx]],
# [x[x_min_idx], y[y_max_idx]],
# [x[x_min_idx], y[y_min_idx]]
# ])
# avg_max_length = (y[y_max_idx]- y[y_min_idx])
# image = mask_array*255.0 # class, proposed, width, height
# im = Image.fromarray(image)
# im = im.convert("L")
# im.save(f"./images_for_check/special{i}_{k}.png") # why only 100??
class_pred = k + 1
poly_results_per_image.append(
(class_pred, class_bbox[l][-1], coords))
imcheck = Image.fromarray(mask_array).convert(
"RGB")
img1 = ImageDraw.Draw(imcheck)
# coords_tuple = [(round(coord[0]), round(coord[1]))for coord in coords]
# img1.polygon(coords_tuple, fill=(128, 0, 0, 50), outline ="blue")
# imcheck.save(f"./images_for_check2/special{i}_{k}.png") # why only 100??
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
if out:
poly_results.append(poly_results_per_image)
for _ in range(batch_size):
prog_bar.update()
if out:
assert len(poly_results) == len(filenames)
if out:
with open('../your_file2.txt', 'w') as f:
for item in sides_list:
f.write("%s\n" % item)
return results, poly_results, filenames
else:
return results
def single_gpu_test_custom(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
# TODO: 딱봐도 알지?
annFile = '/content/gdrive/My Drive/k-fashion/data/annotations/instances_test2017.json'
imagepath = '/content/mmdetection/data/test2017/'
coco = COCO(annFile)
cats = coco.loadCats(coco.getCatIds())
nms = [cat['name'] for cat in cats]
print('COCO categories: \n{}\n'.format(' '.join(nms)))
nms = set([cat['supercategory'] for cat in cats])
print('COCO supercategories: \n{}'.format(' '.join(nms)))
# imageid 만들기
filenames = {}
for image_id in coco.getImgIds():
filename = coco.loadImgs(image_id)[0]['file_name']
filenames[filename] = image_id
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
for f in range(2):
if f == 1:
# torch → numpy for image
npimg = data['img'][0].squeeze(0).permute(1, 2, 0).numpy()
filp_ori = cv2.flip(npimg, 1)
# numpy → torch for image
filp_ori = torch.from_numpy(filp_ori).permute(2, 0,
1).unsqueeze(0)
data['img'][0] = filp_ori
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
# 결과도 좌우 전환
#result_maks = result[0][1]
for i, masks in enumerate(result[0][1]):
if len(masks) == 0: continue
for j, mask in enumerate(masks):
fmask = cv2.flip(mask * 1, 1)
result[0][1][i][j] = fmask.astype(bool)
# imageid 추가
result[0] = result[0] + tuple([
filenames[data['img_metas'][0].data[0][0]['ori_filename']]
])
else:
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
# imageid 추가
result[0] = result[0] + tuple([
filenames[data['img_metas'][0].data[0][0]
['ori_filename']]
])
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results),
imageid)
for bbox_results, mask_results, imageid in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
fps=3,
show_score_thr=0.3):
"""Test model with single gpu.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): Pytorch data loader.
show (bool, optional): If True, visualize the prediction results.
Defaults to False.
out_dir (str, optional): Path of directory to save the
visualization results. Defaults to None.
fps (int, optional): FPS of the output video.
Defaults to 3.
show_score_thr (float, optional): The score threshold of visualization
(Only used in VID for now). Defaults to 0.3.
Returns:
dict[str, list]: The prediction results.
"""
model.eval()
results = defaultdict(list)
dataset = data_loader.dataset
prev_img_meta = None
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = data['img'][0].size(0)
if show or out_dir:
assert batch_size == 1, 'Only support batch_size=1 when testing.'
img_tensor = data['img'][0]
img_meta = data['img_metas'][0].data[0][0]
img = tensor2imgs(img_tensor, **img_meta['img_norm_cfg'])[0]
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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
model.module.show_result(
img_show,
result,
show=show,
out_file=out_file,
score_thr=show_score_thr)
# Whether need to generate a video from images.
# The frame_id == 0 means the model starts processing
# a new video, therefore we can write the previous video.
# There are two corner cases.
# Case 1: prev_img_meta == None means there is no previous video.
# Case 2: i == len(dataset) means processing the last video
need_write_video = (
prev_img_meta is not None and img_meta['frame_id'] == 0
or i == len(dataset))
if out_dir and need_write_video:
prev_img_prefix, prev_img_name = prev_img_meta[
'ori_filename'].rsplit(os.sep, 1)
prev_img_idx, prev_img_type = prev_img_name.split('.')
prev_filename_tmpl = '{:0' + str(
len(prev_img_idx)) + 'd}.' + prev_img_type
prev_img_dirs = f'{out_dir}/{prev_img_prefix}'
prev_img_names = sorted(os.listdir(prev_img_dirs))
prev_start_frame_id = int(prev_img_names[0].split('.')[0])
prev_end_frame_id = int(prev_img_names[-1].split('.')[0])
mmcv.frames2video(
prev_img_dirs,
f'{prev_img_dirs}/out_video.mp4',
fps=fps,
fourcc='mp4v',
filename_tmpl=prev_filename_tmpl,
start=prev_start_frame_id,
end=prev_end_frame_id,
show_progress=False)
prev_img_meta = img_meta
for key in result:
if 'mask' in key:
result[key] = encode_mask_results(result[key])
for k, v in result.items():
results[k].append(v)
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test_crop_img(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3):
print('clw: using single_gpu_test_crop_img() !!')
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(
data_loader): # data['img][0]: tensor (1, 3, 6016, 8192)
img_h = data['img'][0].shape[2]
img_w = data['img'][0].shape[3]
with torch.no_grad():
# 如果是 4096x3500,直接原图预测
if img_h <= 3500 and img_w <= 4096:
#if img_h <= 10000 and img_w <= 10000:
result = model(return_loss=False, rescale=True, **data)
else:
# 否则切图, 4 nums
##############################
overlap_h = 272
overlap_w = 256
crop_h = round(
(img_h + overlap_h) /
2) # clw note: the size can be divided by 32 is better
crop_w = round((img_w + overlap_w) / 2)
#crop_h = 800
#crop_w = 1344
# crop_w = 1333
#step_h = int(0.8 * crop_h)
#step_w = int(0.8 * crop_w)
step_h = crop_h - overlap_h
step_w = crop_w - overlap_w
nms_iou_thr = model.module.test_cfg['rcnn']['nms'][
'iou_threshold']
results_crop = [[] for _ in range(len(model.module.CLASSES))]
data['img_metas'][0].data[0][0]['ori_shape'] = (crop_h, crop_w)
data['img_metas'][0].data[0][0]['img_shape'] = (crop_h, crop_w)
data['img_metas'][0].data[0][0]['pad_shape'] = (crop_h, crop_w)
img_tensor_orig = data['img'][0].clone()
for start_h in range(0, img_h - crop_h + 1,
step_h): # imgsz is crop step here,
if start_h + crop_h > img_h: # 如果最后剩下的不到imgsz,则step少一些,保证切的图尺寸不变
start_h = img_h - crop_h
for start_w in range(0, img_w - crop_w + 1, step_w):
if start_w + crop_w > img_w: # 如果最后剩下的不到imgsz,则step少一些,保证切的图尺寸不变
start_w = img_w - crop_w
# crop
print(start_h, start_w)
data['img'][0] = img_tensor_orig[:, :,
start_h:start_h +
crop_h,
start_w:start_w +
crop_w]
result = model(
return_loss=False, rescale=True, **data
) # result[0]: model.module.CLASSES 个list,每个里面装着(n, 5) ndarray
#result = model(return_loss=False, rescale=False, **data) # clw modify
for idx, item in enumerate(result[0]):
for row in item:
#print('boxw:', row[2] - row[0], 'boxh:', row[3] - row[1] )
if row[2] - row[0] == 0 or row[3] - row[1] == 0:
print(
'==================================================================='
)
continue
row[[0, 2]] += start_w
row[[1, 3]] += start_h
results_crop[idx].append(row)
results_afternms = []
for idx, res in enumerate(results_crop):
if len(res) == 0:
results_afternms.append(np.array(
[])) # clw note: it's really important!!
continue
else:
prediction = torch.tensor(res)
boxes, scores = prediction[:, :
4], prediction[:,
4] # boxes (offset by class), scores
i = torchvision.ops.boxes.nms(boxes, scores,
nms_iou_thr)
results_afternms.append(prediction[i].numpy())
result = [results_afternms]
##############################
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
return results
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show \
or args.show_dir, \
('Please specify at least one operation (save/eval/format/show the '
'results / save the results) with the argument "--out", "--eval"'
', "--format-only", "--show" or "--show-dir"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# import modules from string list.
if cfg.get('custom_imports', None):
from mmcv.utils import import_modules_from_strings
import_modules_from_strings(**cfg['custom_imports'])
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
if cfg.model.get('neck'):
if isinstance(cfg.model.neck, list):
for neck_cfg in cfg.model.neck:
if neck_cfg.get('rfp_backbone'):
if neck_cfg.rfp_backbone.get('pretrained'):
neck_cfg.rfp_backbone.pretrained = None
elif cfg.model.neck.get('rfp_backbone'):
if cfg.model.neck.rfp_backbone.get('pretrained'):
cfg.model.neck.rfp_backbone.pretrained = None
# in case the test dataset is concatenated
if isinstance(cfg.data.test, dict):
cfg.data.test.test_mode = True
elif isinstance(cfg.data.test, list):
for ds_cfg in cfg.data.test:
ds_cfg.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
if samples_per_gpu > 1:
# Replace 'ImageToTensor' to 'DefaultFormatBundle'
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.fuse_conv_bn:
model = fuse_conv_bn(model)
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
model = MMDataParallel(model, device_ids=[0])
model.eval()
dataset = data_loader.dataset
iou = []
cls_conf = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
show = args.show
out_dir = args.show_dir
show_score_thr = args.show_score_thr
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
# results.extend(result)
#
# for _ in range(batch_size):
# prog_bar.update()
# return results
with torch.no_grad():
for index in range(len(dataset_val)):
retinanet.eval()
data = dataset_val[index]
scale = data['scale']
# run network
classification, transformed_anchors = retinanet(
data['img'].permute(2, 0, 1).cuda().float().unsqueeze(dim=0))
classification = classification.squeeze(0)
transformed_anchors = transformed_anchors.squeeze(0)
scores = torch.max(classification, dim=1, keepdim=True)[1] # 类别的索引
scores_key = torch.unique(scores) # 一张图片中预测结果包含哪些类别的目标
if len(scores_key) == 1 and scores_key[0] == 0.0:
continue
target = data['annot'].cuda()
target_cls = torch.unique(target[:, 4]) # 一张图片中真实结果包含哪些类别的目标
for idx in range(scores_key.shape[0]):
cls_item = scores_key[idx].item()
cls_conf_tmp = classification[(
scores[:, :] == cls_item).squeeze()][:, cls_item]
pred_reg = transformed_anchors[(
scores[:, :] == cls_item).squeeze()]
if cls_item in target_cls:
target_tmp = target[:, :4][target[:, 4] == cls_item]
iou_tmp = torch.max(calc_iou(pred_reg, target_tmp),
dim=1,
keepdim=True)[0].squeeze(1)
# iou.extend(iou_tmp.cpu().numpy())
# cls_conf.extend(cls_conf_tmp.cpu().numpy())
else:
iou_tmp = torch.zeros_like(cls_conf_tmp)
iou.extend(iou_tmp.cpu().numpy())
cls_conf.extend(cls_conf_tmp.cpu().numpy())
if index == 500:
break
# cls_count = 0
# iou_count = 0
print("检测框个数： ", len(cls_conf))
# for i in range(len(cls_conf)):
# print(cls_conf[i])
# if cls_conf[i] >= 0.5:
# cls_count += 1
# if iou[i] >= 0.5:
# iou_count += 1
# print("分类置信度>0.5:", cls_count)
# print("定位置信度(IoU)>0.5:", iou_count)
print("分类置信度>0.5:", np.sum(np.greater_equal(cls_conf, 0.5)))
print("定位置信度(IoU)>0.5:", np.sum(np.greater_equal(iou, 0.9)))
plt.scatter(iou, cls_conf, marker='x', color='red', s=1)
plt.savefig('iou_result/scatter_335.jpg')
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3,
output_guangdong=False):
model.eval()
results = []
names = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
# import pdb
# pdb.set_trace()
batch_size = len(result)
if show or out_dir:
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 = osp.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)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
if output_guangdong:
names.extend([
meta['ori_filename'] for meta in data['img_metas'][0].data[0]
])
results.extend(result)
# import pdb
# pdb.set_trace()
# if len(results)>=100:
# return [results,names]
# return results
for _ in range(batch_size):
prog_bar.update()
return results if not output_guangdong else [results, names]
def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
show_score_thr=0.3,
export_file=None
):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
annotation_list = []
image_id = 1
for i, data in enumerate(data_loader):
if export_file:
image_id = i + 1
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
batch_size = len(result)
if show or out_dir:
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 = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
if export_file:
#################wgp add##################
if isinstance(result[i], tuple):
bbox_result, segm_result = result[i]
if isinstance(segm_result, tuple):
segm_result = segm_result[0] # ms rcnn
else:
bbox_result, segm_result = result[i], None
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
# image_name = out_file.split('/')[-1]
for k in range(len(bboxes)):
category_id = int(labels[k])
x1, y1, x2, y2 = bboxes[k][:4]
score = float(bboxes[k][-1])
if score > show_score_thr:
annotation_item_dict = {'image_id': image_id, 'category_id': category_id + 1,
'bbox': [x1.item(), y1.item(), x2.item() - x1.item(),
y2.item() - y1.item()], 'score': score}
annotation_list.append(annotation_item_dict)
###################################
model.module.show_result(
img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
# encode mask results
if isinstance(result[0], tuple):
result = [(bbox_results, encode_mask_results(mask_results))
for bbox_results, mask_results in result]
results.extend(result)
for _ in range(batch_size):
prog_bar.update()
if export_file:
if out_dir is None:
out_dir = '.'
annotation_json_str = json.dumps(
annotation_list,
ensure_ascii=False,
indent=4,
separators=[',', ': ']
)
with open(out_dir + f'/{export_file}', 'w', encoding='utf-8') as json_file:
json_file.writelines(annotation_json_str)
return results
colors = [(55, 255, 20), (0, 0, 255), (132, 240, 255), (0, 247, 255),
(2, 2, 105)]
filename = [
"paragraph_boxes.jpg", "header_boxes.jpg", "list_boxes.jpg",
"tabel_boxes.jpg", "figure_boxes.jpg"
]
for box in category:
#print (count)
#print(colors[count])
cv2.rectangle(im1, (box[0], box[1]), (box[2], box[3]), colors[count],
2)
cv2.rectangle(im2, (box[0], box[1]), (box[2], box[3]), colors[count],
2)
directory = '/content/drive/My Drive/results'
os.chdir(directory)
name = filename[count]
#print (name)
cv2.imwrite(name, im1)
directory = '/content/drive/My Drive/results'
os.chdir(directory)
result_file = "all_annotations.jpg"
cv2.imwrite(result_file, im2)
encoded_mask_results = encode_mask_results(mask_results)
print("Encoded mask results are == ", encoded_mask_results)
result = bbox_results, encoded_mask_results
results.append(result)
