def after_nms(nms_outs,
img_h,
img_w,
show_lincomb=False,
crop_masks=True,
visual_thre=0,
img_name=None):
if nms_outs is None:
return [torch.Tensor()
] * 4 # Warning, this is 4 copies of the same thing
if visual_thre > 0:
keep = nms_outs['class'] > visual_thre
for k in nms_outs:
if k != 'proto':
nms_outs[k] = nms_outs[k][keep]
if nms_outs['class'].size(0) == 0:
return [torch.Tensor()] * 4
class_ids = nms_outs['class_ids']
boxes = nms_outs['box']
classes = nms_outs['class']
coefs = nms_outs['coef']
# At this points masks is only the coefficients
proto_data = nms_outs['proto']
if show_lincomb:
draw_lincomb(proto_data, coefs, img_name)
masks = torch.sigmoid(torch.matmul(proto_data, coefs.t()))
# Crop masks by boxes
if crop_masks:
masks = crop(masks, boxes)
masks = masks.permute(2, 0, 1).contiguous()
masks = F.interpolate(masks.unsqueeze(0), (img_h, img_w),
mode='bilinear',
align_corners=False).squeeze(0)
# Binarize the masks
masks.gt_(0.5)
boxes[:, 0], boxes[:, 2] = sanitize_coordinates(boxes[:, 0], boxes[:, 2],
img_w)
boxes[:, 1], boxes[:, 3] = sanitize_coordinates(boxes[:, 1], boxes[:, 3],
img_h)
boxes = boxes.long()
return class_ids, classes, boxes, masks
def after_nms(nms_outs, img_h, img_w, cfg=None, img_name=None):
if nms_outs is None:
return [torch.Tensor()] * 4
if cfg and cfg.visual_thre > 0:
keep = nms_outs['class'] >= cfg.visual_thre
for k in nms_outs:
if k != 'proto':
nms_outs[k] = nms_outs[k][keep]
if nms_outs['class'].size(0) == 0:
return [torch.Tensor()] * 4
class_ids = nms_outs['class_ids']
boxes = nms_outs['box']
classes = nms_outs['class']
coefs = nms_outs['coef']
proto_data = nms_outs['proto']
if cfg and cfg.save_lincomb:
draw_lincomb(proto_data, coefs, img_name)
masks = torch.sigmoid(torch.matmul(proto_data, coefs.t()))
if not cfg or not cfg.no_crop: # Crop masks by boxes
masks = crop(masks, boxes)
masks = masks.permute(2, 0, 1).contiguous()
masks = F.interpolate(masks.unsqueeze(0), (img_h, img_w),
mode='bilinear',
align_corners=False).squeeze(0)
masks.gt_(0.5) # Binarize the masks
boxes[:, 0], boxes[:, 2] = sanitize_coordinates(boxes[:, 0], boxes[:, 2],
img_w)
boxes[:, 1], boxes[:, 3] = sanitize_coordinates(boxes[:, 1], boxes[:, 3],
img_h)
boxes = boxes.long()
return class_ids, classes, boxes, masks