def init_detector(config, checkpoint=None, device='cuda:0'):
"""Initialize a detector from config file.
Args:
config (str or :obj:`mmcv.Config`): Config file path or the config
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
Returns:
nn.Module: The constructed detector.
"""
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
'but got {}'.format(type(config)))
config.model.pretrained = None
model = build_detector(config.model, test_cfg=config.test_cfg)
if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint)
if 'CLASSES' in checkpoint['meta']:
# model.CLASSES = checkpoint['meta']['CLASSES']
model.CLASSES = get_classes('own')
else:
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.')
model.CLASSES = get_classes('own')
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()
return model
def init_detector(config, checkpoint=None, device='cuda:0'):
"""Initialize a detector from config file.
Args:
config (str or :obj:`mmcv.Config`): Config file path or the config
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
Returns:
nn.Module: The constructed detector.
"""
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}')
config.model.pretrained = None
model = build_detector(config.model, test_cfg=config.test_cfg)
if checkpoint is not None:
map_loc = 'cpu' if device == 'cpu' else None
checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
warnings.simplefilter('once')
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.')
model.CLASSES = get_classes('coco')
# model = MMDataParallel(model, device_ids=[0]) # 직접 추가한 코드, inference time 체크할 때 사용
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()
return model
def show_result(self,
data,
result,
img_norm_cfg,
frame_index,
show=False,
dataset=None,
score_thr=0.3):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
img_tensor = data['img'][0]
img_metas = data['img_meta'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32) # category_id -1
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
show=show,
class_names=class_names,
wait_time=1,
out_file="output/polyp/" +
str("{:04d}".format(frame_index)) + ".png",
score_thr=score_thr)
def init_detector(config, checkpoint=None, device='cuda:0'):
"""Initialize a detector from config file.
Args:
config (str or :obj:`mmcv.Config`): Config file path or the config
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
Returns:
nn.Module: The constructed detector.
"""
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}')
config.model.pretrained = None
model = build_detector(config.model, test_cfg=config.test_cfg)
if checkpoint is not None:
map_loc = 'cpu' if device == 'cpu' else None
checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = get_classes('coco')
model.eval()
return model
def show_result(img, result, dataset='coco', score_thr=0.3, out_file=None, font_scale=0.5):
img = mmcv.imread(img)
class_names = get_classes(dataset)
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(
0, 256, (1, 3), dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img[mask] = img[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
write_bboxes_to_npy(bboxes, out_file)
mmcv.imshow_det_bboxes(
img.copy(),
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
show=out_file is None,
out_file=out_file,
font_scale=font_scale)
def vis_seg(data, result, img_norm_cfg, data_id, colors, score_thr, save_dir):
img_tensor = data['img'][0]
img_metas = data['img_meta'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
class_names = get_classes('tianchi')
for img, img_meta, cur_result in zip(imgs, img_metas, result):
if cur_result is None:
continue
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
seg_label = cur_result[0]
seg_label = seg_label.cpu().numpy().astype(np.uint8)
cate_label = cur_result[1]
cate_label = cate_label.cpu().numpy()
score = cur_result[2].cpu().numpy()
vis_inds = score > score_thr
seg_label = seg_label[vis_inds]
num_mask = seg_label.shape[0]
cate_label = cate_label[vis_inds]
cate_score = score[vis_inds]
mask_density = []
for idx in range(num_mask):
cur_mask = seg_label[idx, :, :]
cur_mask = mmcv.imresize(cur_mask, (w, h))
cur_mask = (cur_mask > 0.5).astype(np.int32)
mask_density.append(cur_mask.sum())
orders = np.argsort(mask_density)
seg_label = seg_label[orders]
cate_label = cate_label[orders]
cate_score = cate_score[orders]
seg_show = img_show.copy()
for idx in range(num_mask):
idx = -(idx+1)
cur_mask = seg_label[idx, :,:]
cur_mask = mmcv.imresize(cur_mask, (w, h))
cur_mask = (cur_mask > 0.5).astype(np.uint8)
if cur_mask.sum() == 0:
continue
color_mask = np.random.randint(
0, 256, (1, 3), dtype=np.uint8)
cur_mask_bool = cur_mask.astype(np.bool)
seg_show[cur_mask_bool] = img_show[cur_mask_bool] * 0.5 + color_mask * 0.5
cur_cate = cate_label[idx]
cur_score = cate_score[idx]
label_text = class_names[cur_cate]
#label_text += '|{:.02f}'.format(cur_score)
# center
center_y, center_x = ndimage.measurements.center_of_mass(cur_mask)
vis_pos = (max(int(center_x) - 10, 0), int(center_y))
cv2.putText(seg_show, label_text, vis_pos,
cv2.FONT_HERSHEY_COMPLEX, 0.3, (255, 255, 255)) # green
mmcv.imwrite(seg_show, '{}/{}.jpg'.format(save_dir, data_id))
def show_gt_result(self,
data,
bbox_gts,
label_gts,
img_norm_cfg,
mask_gts=None,
dataset=None,
image_dir=None,
show=False,
image_name=None):
bbox_gts = bbox_gts.data[0]
label_gts = label_gts.data[0]
if mask_gts:
mask_gts = mask_gts.data[0]
else:
mask_gts = []
img_tensor = data['img'][0]
img_metas = data['img_meta'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta, bbox_gt, label_gt, mask_gt in zip(
imgs, img_metas, bbox_gts, label_gts, mask_gts):
h, w, _ = img_meta['pad_shape']
img_show = img[:h, :w, :]
bboxes = bbox_gt.cpu().numpy()
# draw segmentation masks
if len(mask_gts) > 0:
segms = mask_gt
for i in range(segms.shape[0]):
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = segms[i, :, :].astype(np.bool)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = label_gt.cpu().numpy() - 1
assert image_dir and image_name, 'image_dir and image_name can not be None'
if not osp.exists(image_dir):
os.mkdir(image_dir)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
text_color='white',
class_names=class_names,
score_thr=0,
show=show,
out_file=osp.join(image_dir, image_name))
def show_result(self,
data,
result,
img_norm_cfg,
dataset='coco',
score_thr=0.3):
img_tensor = data['img'][0]
#img_tensor_t = data['img_t'][0]
img_metas = data['img_meta'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
#imgs_t = tensor2imgs(img_tensor_t,**)
assert len(imgs) == len(img_metas)
if isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, list):
class_names = dataset
else:
raise TypeError('dataset must be a valid dataset name or a list'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr)
def initialize_model():
# Setup a checkpoint file to load
#checkpoint = '/Developer/3DObject/mymmdetection/waymococo_fasterrcnnr101train/epoch_60.pth'
#checkpoint = '/Developer/MyRepo/mymodels/mmmodels/HPCwaymococo_fasterrcnnr101train/epoch_25.pth'
device = 'cuda:0'
global model
global config
global model_dir
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}')
config.model.pretrained = None
config.model.train_cfg = None
config.model.roi_head.bbox_head.num_classes = len(classes)
model = build_detector(config.model, test_cfg=config.get('test_cfg'))
if model_dir is not None:
map_loc = 'cpu' if device == 'cpu' else None
checkpoint = load_checkpoint(model, model_dir, map_location=map_loc)
if 'CLASSES' in checkpoint.get('meta', {}):
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = get_classes('coco')
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()
def show_result(self, data, result, dataset=None, score_thr=0.3):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
img_metas = data['img_meta'][0].data[0]
if 'img_raw' in data:
img_tensor = data['img_raw'][0]
norm_factor = {'mean': [0, 0, 0], 'std': [1, 1, 1], 'to_rgb': True}
filename = os.path.join(
'results',
data['img_meta'][0].data[0][0]['filename'].split('/')[-1])
else:
img_tensor = data['img'][0]
norm_factor = img_metas[0]['img_norm_cfg']
filename = None
imgs = tensor2imgs(img_tensor, **norm_factor)
assert len(imgs) == len(img_metas)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
show=False,
out_file=filename)
def show_result(img,
result,
dataset,
score_thr=0.03,
wait_time=0,
show=True,
out_file=None):
"""Visualize the detection results on the image.
Args:
img (str or np.ndarray): Image filename or loaded image.
result (tuple[list] or list): The detection result, can be either
(bbox, segm) or just bbox.
class_names (list[str] or tuple[str]): A list of class names.
score_thr (float): The threshold to visualize the bboxes and masks.
wait_time (int): Value of waitKey param.
show (bool, optional): Whether to show the image with opencv or not.
out_file (str, optional): If specified, the visualization result will
be written to the out file instead of shown in a window.
Returns:
np.ndarray or None: If neither `show` nor `out_file` is specified, the
visualized image is returned, otherwise None is returned.
"""
# print("score阈值设置为%f"%score_thr)
class_names = get_classes(dataset)
assert isinstance(class_names, (tuple, list))
img = mmcv.imread(img)
img = img.copy()
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0, 256, (1, 3), dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img[mask] = img[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
anno, img = mmcv.imshow_det_bboxes(img,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
show=show,
wait_time=wait_time,
out_file=out_file)
if not (show or out_file):
return anno, img
def show_result(
self,
data,
result,
img_norm_cfg,
dataset="coco",
score_thr=0.3,
):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
img_tensor = data["img"][0]
img_metas = data["img_meta"][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
if isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)) or dataset is None:
class_names = dataset
else:
raise TypeError(
"dataset must be a valid dataset name or a sequence"
" of class names, not {}".format(type(dataset))
)
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta["img_shape"]
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(
0, 256, (1, 3), dtype=np.uint8
)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img_show[mask] = (
img_show[mask] * 0.5 + color_mask * 0.5
)
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(
img_show,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
)
def det_bboxes(img,
bboxes,
labels,
score_thr=0,
bbox_color='green',
text_color='green',
thickness=2,
font_scale=0.5,
show=True,
win_name='',
wait_time=0,
out_file=None):
assert bboxes.ndim == 2
assert labels.ndim == 1
assert bboxes.shape[0] == labels.shape[0]
assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
if score_thr > 0:
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
bbox_color = mmcv.color_val(bbox_color)
text_color = mmcv.color_val(text_color)
ret_str = ''
for bbox, label in zip(bboxes, labels):
bbox_int = bbox.astype(np.int32)
left_top = (bbox_int[0], bbox_int[1])
right_bottom = (bbox_int[2], bbox_int[3])
cv2.rectangle(img,
left_top,
right_bottom,
bbox_color,
thickness=thickness)
label_text = str(label)
dataset_name = 'sign'
class_names = get_classes(dataset_name)
label_text = class_names[
label] if class_names is not None else 'cls {}'.format(label)
w = bbox_int[2] - bbox_int[0]
h = bbox_int[3] - bbox_int[1]
ss = ' ' + str(left_top[0]) + ',' + str(
left_top[1]) + ',' + str(w) + ',' + str(h)
if len(bbox) > 4:
cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_scale, text_color)
ret_str += ss
if show:
mmcv.imshow(img, win_name, wait_time)
if out_file is not None:
mmcv.imwrite(img, out_file)
return ret_str
def show_result(self,
data,
result,
img_norm_cfg,
dataset=None,
score_thr=0.3,
outfile=None):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
img_tensor = data['img'][0]
img_metas = data['img_meta'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_norm_cfg)
assert len(imgs) == len(img_metas)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
# solve UnboundLocalError: local variable '_mlvl_bboxes' referenced before assignment
# h, w, _ = img_meta['img_shape']
h, w, _ = img_meta['ori_shape']
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
out_file=outfile)
def show_result(self, data, result, dataset=None, score_thr=0.3):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
img_tensor = data['img'].data
print(img_tensor.shape)
img_metas = data['img_meta'].data[0]
print(img_metas)
imgs = tensor2imgs(
img_tensor, **{
'mean': np.array([0.5, 0.5, 0.5]),
'std': np.array([255., 255., 255.]),
'to_rgb': False
})
assert len(imgs) == len(img_metas)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
print(segms)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
print('amsk shape', mask.shape)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(img_show,
bboxes,
labels,
class_names=class_names,
score_thr=score_thr)
def init_detector(opts, device='cuda:0'):
config = mmcv.Config.fromfile(opts.config)
new_config = 'train_pipeline' in config or 'test_pipeline' in config
if new_config:
assert opts.in_scale is not None, "New config does not support default input scale"
# simulate old config
config.data.test.img_scale = 1
config.data.test.size_divisor = 32
if opts.in_scale is not None:
if 'ssd' in basename(opts.config):
# SSD
if opts.in_scale <= 0.2:
# too small leads to some issues
l = round(1920 * opts.in_scale)
config.data.test.img_scale = (l, l)
config.data.test.resize_keep_ratio = False
else:
config.data.test.img_scale = opts.in_scale
config.data.test.resize_keep_ratio = True
# My change
elif config.model.type == 'YoloV3':
#elif config.model.type == 'YoloV3' or config.model.type == 'YoloV3UA' :
config.data.test.in_scale = opts.in_scale
config.data.test.resize_keep_ratio = True
else:
config.data.test.img_scale = opts.in_scale
config.data.test.resize_keep_ratio = True
if opts.no_mask:
if 'mask_head' in config.model:
config.model['mask_head'] = None
config.model.pretrained = None
model = build_detector(config.model, test_cfg=config.test_cfg)
if opts.weights is not None:
if config.model.type == 'YoloV3':
if opts.weights.endswith(".weights"):
# Load darknet weights
# Check if the weights using are the same. '/data2/mengtial/ModelZoo/yolo/yolov3.weights'
model.load_darknet_weights(opts.weights)
else:
# Load checkpoint weights
model.load_state_dict(torch.load(opts.weights))
elif config.model.type == 'YoloV3UA':
model.load_weights(opts.weights)
else:
weights = load_checkpoint(model, opts.weights)
if 'CLASSES' in weights['meta']:
model.CLASSES = weights['meta']['CLASSES']
else:
model.CLASSES = get_classes('coco')
model.cfg = config
model.to(device)
model.eval()
return model
def init_detector(opts, device='cuda:0'):
config = mmcv.Config.fromfile(opts.config)
new_config = 'train_pipeline' in config or 'test_pipeline' in config
if new_config:
# simulate old config
if opts.in_scale is None:
print('Warning: using new config and fixing size_divisor to 32')
config.data.test.img_scale = config.test_pipeline[1]['img_scale']
else:
config.data.test.img_scale = 1
config.data.test.size_divisor = 32
if opts.in_scale is not None:
if 'ssd' in basename(opts.config):
# SSD
if opts.in_scale <= 0.2:
# too small leads to some issues
l = round(1920 * opts.in_scale)
config.data.test.img_scale = (l, l)
config.data.test.resize_keep_ratio = False
else:
config.data.test.img_scale = opts.in_scale
config.data.test.resize_keep_ratio = True
else:
config.data.test.img_scale = opts.in_scale
config.data.test.resize_keep_ratio = True
if opts.no_mask:
if 'roi_head' in config.model and 'mask_head' in config.model[
'roi_head']:
config.model['roi_head']['mask_head'] = None
if 'zoom_crop' in opts and opts.zoom_crop:
config.data.test.zoom_crop = {
'h': opts.zoom_crop_h,
'y': opts.zoom_crop_y,
}
else:
config.data.test.zoom_crop = None
config.model.pretrained = None
if 'action_head' in config.model:
config.model['action_head_weights'] = opts.action_head_weights
model = build_detector(config.model, test_cfg=config.test_cfg)
map_loc = 'cpu' if device == 'cpu' else None
checkpoint = load_checkpoint(model, opts.weights, map_location=map_loc)
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
warnings.simplefilter('once')
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.')
model.CLASSES = get_classes('coco')
model.cfg = config
model.to(device)
model.eval()
return model
def show_and_save_result(img, result, out_dir, dataset="coco", score_thr=0.3):
class_names = get_classes(dataset)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
index = img.rfind("/")
mmcv.imshow_det_bboxes(img, bboxes, labels, class_names, score_thr, show=True, out_file=out_dir+img[index+1:])
def show_result(self, data, result, dataset=None, score_thr=0.3):
if isinstance(result, tuple):
bbox_result, segm_result = result
else:
bbox_result, segm_result = result, None
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)
if dataset is None:
class_names = self.CLASSES
elif isinstance(dataset, str):
class_names = get_classes(dataset)
elif isinstance(dataset, (list, tuple)):
class_names = dataset
else:
raise TypeError(
'dataset must be a valid dataset name or a sequence'
' of class names, not {}'.format(type(dataset)))
for img, img_meta in zip(imgs, img_metas):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
bboxes = np.vstack(bbox_result)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
for i in inds:
color_mask = np.random.randint(0,
256, (1, 3),
dtype=np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool)
img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
print("img_meta: {}".format(img_meta))
mmcv.imshow_det_bboxes(
img_show,
bboxes,
labels,
# class_names=class_names,
class_names=[str(i) for i in range(len(class_names))],
score_thr=score_thr,
gt_bboxes=data['gt_bboxes'][0][0].numpy(),
gt_labels=(data['gt_labels'][0][0] - 1).numpy())
def parse_pkldets(input_pkl, img_dir, outdir, score_thr=0.3, dataset='cityscapes'):
dets = mmcv.load(input_pkl)
num_imgs = len(dets)
# num_class = 8
VOC_CLASSES = get_classes(dataset)
print(VOC_CLASSES)
img_names_list = os.listdir(img_dir)
for k in range(num_imgs):
print(k)
imgHeight, imgWidth = 1024, 2048
imgName = img_names_list[k]
# print(dets[k])
bbox_result, segm_result = dets[k]
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(bbox_result)
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
# labs, confs, polygons = [], [], []
ftxt = open(outdir + imgName[:-4] + '.txt', 'w')
# print(inds)
for i in inds:
img = np.zeros((imgHeight, imgWidth), np.uint8)
mask = maskUtils.decode(segms[i]).astype(np.bool) # mask predicted.
# contour, hier = cv2.findContours(mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
# try:
# contour = np.reshape(contour[0], [-1, 2])
# except:
# print(k, i, contour, 'Segmentation Error!')
# continue
# polygon = []
# for j in range(contour.shape[0]):
# polygon.append([int(contour[j, 0]), int(contour[j, 1])])
# polygons.append(polygon) # append polygon of a instance.
img[mask] = 255 # mask as white-color.
cv2.imwrite(outdir + imgName[:-4] + str(i) + '.png', img)
# confs.append(round(float(bboxes[i, -1]), 2)) # conf predicted.
# labs.append(VOC_CLASSES[labels[i]]) # labels predicted.
# print(imgName[:-4]+str(i)+'.png')
# print(str(labelDic[VOC_CLASSES[labels[i]]]))
ftxt.write(imgName[:-4]+str(i)+'.png'+' '+str(labelDic[VOC_CLASSES[labels[i]]])+' '+str(round(float(bboxes[i, -1]), 2))+'\n')
def show_result(img, result, dataset='coco', score_thr=0.3):
class_names = get_classes(dataset)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
img = mmcv.imread(img)
mmcv.imshow_det_bboxes(img.copy(),
bboxes,
labels,
class_names=class_names,
score_thr=score_thr)
def run_on_onnxruntime():
config_file = '../configs/solov2/solov2_light_448_r34_fpn_8gpu_3x.py'
onnx_file = 'weights/SOLOv2_light_R34.onnx'
input_names = ['input']
# output_names = ['C0', 'C1', 'C2', 'C3', 'C4']
# output_names = ['cate_pred_0', 'cate_pred_1', 'cate_pred_2', 'cate_pred_3', 'cate_pred_4',
# 'kernel_pred_0', 'kernel_pred_1', 'kernel_pred_2', 'kernel_pred_3', 'kernel_pred_4',
# 'seg_pred'] # Origin
output_names = ['cate_pred', 'kernel_pred', 'seg_pred'] # add permute & concate
if isinstance(config_file, str):
cfg = mmcv.Config.fromfile(config_file)
elif not isinstance(config_file, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
'but got {}'.format(type(config_file)))
# 1. Preprocess
# input demo img size 427x640 --> resized 448x671 --> pad 448x672
img = 'images/demo.jpg'
# 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)
# 2. Run inference on onnxruntime
print("Load onnx model from {}.".format(onnx_file))
sess = rt.InferenceSession(onnx_file)
tic = cv2.getTickCount()
onnx_output = sess.run(output_names, {input_names[0]: data['img'][0].unsqueeze(0).cpu().numpy()})
print('-----> onnxruntime inference time: {}ms'.format((cv2.getTickCount() - tic) * 1000/ cv2.getTickFrequency()))
# 3. Get seg
# 调用pytorch定义的获取分割图和matrix nms 以及后处理
from mmdet.models.anchor_heads.solov2_head import SOLOv2Head
solov2_head = SOLOv2Head(num_classes=81,
in_channels=256,
num_grids=[40, 36, 24, 16, 12],
strides=[8, 8, 16, 32, 32],
ins_out_channels = 128,
loss_ins=cfg.model.bbox_head.loss_ins,
loss_cate=cfg.model.bbox_head.loss_cate)
cate_preds = [torch.from_numpy(x) for x in onnx_output[:1]]
kernel_preds = [torch.from_numpy(x) for x in onnx_output[1:2]]
seg_pred = torch.from_numpy(onnx_output[2])
result = solov2_head.get_seg(cate_preds, kernel_preds, seg_pred, [data['img_meta'][0].data], cfg.test_cfg, rescale=True)
show_result_ins(img, result, get_classes('coco'), score_thr=0.25, out_file="images/demo_out_onnxrt_solov2.jpg")
print('Script done!')
def nms_after_det(bboxes_list,
labels_list,
is_pretrained=False,
class_names=get_classes('coco'),
is_rcnn=True):
'''
non maximum suppress on detection result boxes
input bboxes list and labels list, return ndarray of nms result
result format: det_bboxes: [[x1, y1, x2, y2, score],...] det_labels: [0 0 0 0 1 1 1 2 2 ...]
'''
# read config file
if is_rcnn:
cfg_path = FASTER_RCNN_CONFIG
cfg = mmcv.Config.fromfile(cfg_path)
if not is_pretrained:
cfg.model.pretrained = None
# NMS
multi_bboxes = []
multi_scores = []
for i, bbox in enumerate(bboxes_list):
# only show vehicles
# if 2 <= (labels_list[i] + 1) <= 8: # vehicles
if 0 <= labels_list[
i] <= 7: # choose what to keep, now keep person and all vehicles
multi_bboxes.append(bbox[0:4])
# temp = [0 for _ in range(len(class_names))]
temp = [0 for _ in range(3)]
temp[labels_list[i] + 1] = bbox[4]
multi_scores.append(temp)
# if result is null
if not multi_scores:
return np.array([]), np.array([])
if is_rcnn:
det_bboxes, det_labels = multiclass_nms(
torch.from_numpy(np.array(multi_bboxes).astype(np.float32)),
torch.from_numpy(np.array(multi_scores).astype(np.float32)),
cfg.model.test_cfg.rcnn.score_thr, cfg.model.test_cfg.rcnn.nms,
cfg.model.test_cfg.rcnn.max_per_img)
else:
det_bboxes, det_labels = multiclass_nms(
torch.from_numpy(np.array(multi_bboxes).astype(np.float32)),
torch.from_numpy(np.array(multi_scores).astype(np.float32)),
cfg.test_cfg.score_thr, cfg.test_cfg.nms, cfg.test_cfg.max_per_img)
return det_bboxes.numpy(), det_labels.numpy()
def display_result_3d(img, result, dataset='coco', score_thr=0.3):
img_np = np.load(img)
class_names = get_classes(dataset)
bbox_result = result
bboxes = np.vstack(bbox_result)
bboxes_placeholders = [[] for i in range(0, 160)]
for bbox in bboxes:
for z_index in range(int(np.floor(bbox[4])), int(np.ceil(bbox[5])+ 1)):
bboxes_placeholders[z_index].append([bbox[0], bbox[1], bbox[2], bbox[3], bbox[6]])
for index, boxes in enumerate(bboxes_placeholders):
if len(boxes) > 0:
for box in boxes:
if box[4] > score_thr:
print('slice {} score {}'.format(index, box[4]))
def show_result(img, result, dataset='coco', score_thr=0.5, show=True):
class_names = get_classes(dataset)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
# print(result)
# print(len(result))
labels = np.concatenate(labels)
bboxes = np.vstack(result)
img = mmcv.imread(img)
anno, img = mmcv.imshow_det_bboxes(img.copy(),
bboxes,
labels,
show=show,
class_names=class_names,
score_thr=score_thr)
return anno, img
def save_result(result, dataset='coco', score_thr=0.3):
"""Return list dict [{x1,x2,y1,y2,classe,score},...]
Args:
results:
score_thr (float): Minimum score of bboxes to be shown.
"""
class_names = get_classes(dataset)
#class_names = ["person", "plate"]
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
return det_bboxes(bboxes,
labels,
class_names=class_names,
score_thr=score_thr)
def save_result(result, class_to_keep=[], dataset='coco', score_thr=0.3):
"""Return list dict [{x1,x2,y1,y2,classe,score},...]
Args:
results:
class_to_keep (list[str]): Classes to keep (cars, trucks...)
score_thr (float): Minimum score of bboxes to be shown.
"""
class_names = get_classes(dataset)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
return det_bboxes(bboxes,
labels,
class_names=class_names,
class_to_keep=class_to_keep,
score_thr=score_thr)
def init_detector(config, checkpoint=None, device='cuda:0'):
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
'but got {}'.format(type(config)))
config.model.pretrained = None
model = build_recognizer(config.model, test_cfg=config.test_cfg)
if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint)
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.')
model.CLASSES = get_classes('coco')
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()
return model
def save_result(img, result, out_file, dataset='coco', score_thr=0.8):
class_names = get_classes(dataset)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
bboxes = np.vstack(result)
img = mmcv.imread(img)
mmcv.imshow_det_bboxes(img.copy(),
bboxes,
labels,
class_names=class_names,
score_thr=score_thr,
bbox_color='red',
text_color='white',
thickness=2,
font_scale=0.8,
show=False,
out_file=out_file)
def mymm2d_init_detector(config, checkpoint=None, device='cuda:0'):
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}')
config.model.pretrained = None
config.model.train_cfg = None
model = build_detector(config.model, test_cfg=config.get('test_cfg'))
if checkpoint is not None:
map_loc = 'cpu' if device == 'cpu' else None
checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
if 'CLASSES' in checkpoint.get('meta', {}):
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = get_classes('coco')
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()
return model