def poly_nms(self, dets, labels, not_nms_cls=None):
if not_nms_cls is not None:
not_nms_ids = [
np.where(labels == self.class_name.index(cls))[0]
for cls in not_nms_cls
]
nms_ids = [
np.where(labels != self.class_name.index(cls))[0]
for cls in not_nms_cls
]
not_nms_ids = np.hstack(not_nms_ids)
nms_ids = np.hstack(nms_ids)
nms_dets = dets[nms_ids]
ids = poly_nms.poly_gpu_nms(nms_dets, 0.15)
nms_ids = nms_ids[ids]
new_ids = np.hstack([nms_ids, not_nms_ids])
else:
new_ids = poly_nms.poly_gpu_nms(dets, 0.15)
return new_ids
def vis_dets(self, img_name, result):
dets = result['results']
img = cvtools.imread(img_name)
for cls_id, det_cls in dets.items():
det_cls = det_cls[det_cls[:, -1] > 0.5]
if len(det_cls) == 0: continue
ids = poly_nms.poly_gpu_nms(det_cls, 0.15)
det_cls = det_cls[ids]
if cls_id == 15:
img = cvtools.draw_boxes_texts(img, det_cls[:, :-1],
line_width=2,
box_format="polygon")
img_name = osp.basename(img_name)
to_file = osp.join('/code/CenterNet/exp/ctdet/vis', img_name)
cvtools.imwrite(img, to_file)
def vis_dets(ret, img_id, dataset, save_by_cat=False, split='val'):
if split != 'test':
img_info = dataset.coco.loadImgs([img_id])[0]
ann_ids = dataset.coco.getAnnIds(imgIds=[img_id])
anns = dataset.coco.loadAnns(ids=ann_ids)
file_name = img_info['file_name']
else:
img_info = dataset.images(img_id)
file_name = img_info['filename']
colors = []
img = read_dota_image(img_info, split)
for cls_id, det_cls in ret.items():
if len(det_cls) == 0: continue
det_cls = det_cls[det_cls[:, -1] > 0.1]
if len(det_cls) == 0: continue
ids = poly_nms.poly_gpu_nms(det_cls.astype(np.float32), 0.15)
det_cls = det_cls[ids]
text = [
dataset.class_name[cls_id] + str(round(score, 2))
for score in det_cls[..., -1]
]
img = cvtools.draw_boxes_texts(img,
det_cls[:, :-1],
texts=None,
line_width=2,
colors=[dataset.voc_color[cls_id - 1]] *
len(det_cls),
box_format="polygon")
crop_str = list(map(str, img_info['crop']))
filename = osp.splitext(file_name)[0]
suffix = osp.splitext(file_name)[1]
save_img_name = osp.join('_'.join([filename] + crop_str) + suffix)
if save_by_cat and split != 'test':
cats = {ann['category_id'] for ann in anns}
for cat in cats:
cat_name = dataset.coco.cats[cat]['name']
file = osp.join(opt.debug_dir, cat_name, save_img_name)
cvtools.imwrite(img, file)
else:
file = osp.join(opt.debug_dir, save_img_name)
cvtools.imwrite(img, file)
def _nms(dets):
return poly_gpu_nms(dets, thresh, device_id)