def __init__(
self,
pre_nms_top_n,
post_nms_top_n,
nms_thresh,
min_size,
box_coder=None,
fpn_post_nms_top_n=None,
fpn_post_nms_per_batch=True,
):
"""
Arguments:
pre_nms_top_n (int)
post_nms_top_n (int)
nms_thresh (float)
min_size (int)
box_coder (BoxCoder)
fpn_post_nms_top_n (int)
"""
super(RPNPostProcessor, self).__init__()
self.pre_nms_top_n = pre_nms_top_n
self.post_nms_top_n = post_nms_top_n
self.nms_thresh = nms_thresh
self.min_size = min_size
if box_coder is None:
box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))
self.box_coder = box_coder
if fpn_post_nms_top_n is None:
fpn_post_nms_top_n = post_nms_top_n
self.fpn_post_nms_top_n = fpn_post_nms_top_n
self.fpn_post_nms_per_batch = fpn_post_nms_per_batch
def __init__(
self,
pre_nms_thresh,
pre_nms_top_n,
nms_thresh,
fpn_post_nms_top_n,
min_size,
num_classes,
box_coder=None,
):
"""
Arguments:
pre_nms_thresh (float)
pre_nms_top_n (int)
nms_thresh (float)
fpn_post_nms_top_n (int)
min_size (int)
num_classes (int)
box_coder (BoxCoder)
"""
super(RetinaNetPostProcessor, self).__init__(pre_nms_thresh, 0,
nms_thresh, min_size)
self.pre_nms_thresh = pre_nms_thresh
self.pre_nms_top_n = pre_nms_top_n
self.nms_thresh = nms_thresh
self.fpn_post_nms_top_n = fpn_post_nms_top_n
self.min_size = min_size
self.num_classes = num_classes
if box_coder is None:
box_coder = BoxCoder(weights=(10., 10., 5., 5.))
self.box_coder = box_coder
def __init__(self, cfg, in_channels):
super(RPNModule, self).__init__()
self.cfg = cfg.clone()
anchor_generator = make_anchor_generator(cfg)
rpn_head = registry.RPN_HEADS[cfg.MODEL.RPN.RPN_HEAD]
head = rpn_head(cfg, in_channels,
anchor_generator.num_anchors_per_location()[0])
rpn_box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))
box_selector_train = make_rpn_postprocessor(cfg,
rpn_box_coder,
is_train=True)
box_selector_test = make_rpn_postprocessor(cfg,
rpn_box_coder,
is_train=False)
loss_evaluator = make_rpn_loss_evaluator(cfg, rpn_box_coder)
self.anchor_generator = anchor_generator
self.head = head
self.box_selector_train = box_selector_train
self.box_selector_test = box_selector_test
self.loss_evaluator = loss_evaluator
def __init__(
self,
score_thresh=0.05,
nms=0.5,
detections_per_img=100,
box_coder=None,
cls_agnostic_bbox_reg=False,
bbox_aug_enabled=False
):
"""
Arguments:
score_thresh (float)
nms (float)
detections_per_img (int)
box_coder (BoxCoder)
"""
super(PostProcessor, self).__init__()
self.score_thresh = score_thresh
self.nms = nms
self.detections_per_img = detections_per_img
if box_coder is None:
box_coder = BoxCoder(weights=(10., 10., 5., 5.))
self.box_coder = box_coder
self.cls_agnostic_bbox_reg = cls_agnostic_bbox_reg
self.bbox_aug_enabled = bbox_aug_enabled
def __init__(self, cfg):
self.proposal_matcher = Matcher(
cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD,
cfg.MODEL.ROI_HEADS.BG_IOU_THRESHOLD,
allow_low_quality_matches=False,
)
self.fg_bg_sampler = BalancedPositiveNegativeSampler(
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE,
cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION)
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
self.cls_agnostic_bbox_reg = cfg.MODEL.CLS_AGNOSTIC_BBOX_REG
self.sampler_type = cfg.MODEL.ROI_WEAK_HEAD.PARTIAL_LABELS
def __init__(self, cfg, in_channels):
super(RetinaNetModule, self).__init__()
self.cfg = cfg.clone()
anchor_generator = make_anchor_generator_retinanet(cfg)
head = RetinaNetHead(cfg, in_channels)
box_coder = BoxCoder(weights=(10., 10., 5., 5.))
box_selector_test = make_retinanet_postprocessor(cfg,
box_coder,
is_train=False)
loss_evaluator = make_retinanet_loss_evaluator(cfg, box_coder)
self.anchor_generator = anchor_generator
self.head = head
self.box_selector_test = box_selector_test
self.loss_evaluator = loss_evaluator
def make_roi_box_loss_evaluator(cfg):
matcher = Matcher(
cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD,
cfg.MODEL.ROI_HEADS.BG_IOU_THRESHOLD,
allow_low_quality_matches=False,
)
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
box_coder = BoxCoder(weights=bbox_reg_weights)
fg_bg_sampler = BalancedPositiveNegativeSampler(
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE,
cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION)
cls_agnostic_bbox_reg = cfg.MODEL.CLS_AGNOSTIC_BBOX_REG
loss_evaluator = FastRCNNLossComputation(matcher, fg_bg_sampler, box_coder,
cls_agnostic_bbox_reg)
return loss_evaluator
def make_roi_box_post_processor(cfg):
use_fpn = cfg.MODEL.ROI_HEADS.USE_FPN
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
box_coder = BoxCoder(weights=bbox_reg_weights)
score_thresh = cfg.MODEL.ROI_HEADS.SCORE_THRESH
nms_thresh = cfg.MODEL.ROI_HEADS.NMS
detections_per_img = cfg.MODEL.ROI_HEADS.DETECTIONS_PER_IMG
cls_agnostic_bbox_reg = cfg.MODEL.CLS_AGNOSTIC_BBOX_REG
bbox_aug_enabled = cfg.TEST.BBOX_AUG.ENABLED
postprocessor = PostProcessor(
score_thresh,
nms_thresh,
detections_per_img,
box_coder,
cls_agnostic_bbox_reg,
bbox_aug_enabled
)
return postprocessor
class ROI_sampler(object):
def __init__(self, cfg):
self.proposal_matcher = Matcher(
cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD,
cfg.MODEL.ROI_HEADS.BG_IOU_THRESHOLD,
allow_low_quality_matches=False,
)
self.fg_bg_sampler = BalancedPositiveNegativeSampler(
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE,
cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION)
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
self.cls_agnostic_bbox_reg = cfg.MODEL.CLS_AGNOSTIC_BBOX_REG
self.sampler_type = cfg.MODEL.ROI_WEAK_HEAD.PARTIAL_LABELS
def match_targets_to_proposals(self, proposal, target):
match_quality_matrix = boxlist_iou(target, proposal)
matched_idxs = self.proposal_matcher(match_quality_matrix)
# Fast RCNN only need "labels" field for selecting the targets
target = target.copy_with_fields("labels")
# get the targets corresponding GT for each proposal
matched_targets = target[matched_idxs.clamp(min=0)]
matched_targets.add_field("matched_idxs", matched_idxs)
return matched_targets
def prepare_targets(self, proposals, targets):
labels = []
regression_targets = []
for proposals_per_image, targets_per_image in zip(proposals, targets):
matched_targets = self.match_targets_to_proposals(
proposals_per_image, targets_per_image)
matched_idxs = matched_targets.get_field("matched_idxs")
labels_per_image = matched_targets.get_field("labels")
labels_per_image = labels_per_image.to(dtype=torch.int64)
# Label background (below the low threshold)
bg_inds = matched_idxs == Matcher.BELOW_LOW_THRESHOLD
labels_per_image[bg_inds] = 0
# Label ignore proposals (between low and high thresholds)
ignore_inds = matched_idxs == Matcher.BETWEEN_THRESHOLDS
labels_per_image[ignore_inds] = -1 # -1 is ignored by sampler
# compute regression targets
regression_targets_per_image = self.box_coder.encode(
matched_targets.bbox, proposals_per_image.bbox)
labels.append(labels_per_image)
regression_targets.append(regression_targets_per_image)
return labels, regression_targets
def subsample_box(self, proposals, targets):
labels, regression_targets = self.prepare_targets(proposals, targets)
sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
proposals = list(proposals)
# add corresponding label and regression_targets information to the bounding boxes
for labels_per_image, regression_targets_per_image, proposals_per_image in zip(
labels, regression_targets, proposals):
proposals_per_image.add_field("labels", labels_per_image)
proposals_per_image.add_field("regression_targets",
regression_targets_per_image)
# distributed sampled proposals, that were obtained on all feature maps
# concatenated via the fg_bg_sampler, into individual feature map levels
for img_idx, (pos_inds_img, neg_inds_img) in enumerate(
zip(sampled_pos_inds, sampled_neg_inds)):
img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img,
as_tuple=False).squeeze(1)
proposals_per_image = proposals[img_idx][img_sampled_inds]
proposals[img_idx] = proposals_per_image
return proposals
def match_targets_to_proposals_point(self, proposal, target):
clicks = target.get_field('click').keypoints
clicks_tiled = torch.unsqueeze(torch.cat((clicks, clicks), dim=1),
dim=1)
num_obj = clicks.shape[0]
box_repeat = torch.cat(
[proposal.bbox.unsqueeze(0) for _ in range(num_obj)], dim=0)
diff = clicks_tiled - box_repeat
matched_ids = (diff[:, :, 0] > 0) * (diff[:, :, 1] > 0) * (
diff[:, :, 2] < 0) * (diff[:, :, 3] < 0)
# Pytorch bug: https://discuss.pytorch.org/t/runtimeerror-argmax-cuda-not-implemented-for-bool/58808
matched_idxs = matched_ids.float().argmax(0)
matched_idxs[matched_ids.sum(0) == 0] = -1
# print((matched_idxs>0).sum(0))
# Fast RCNN only need "labels" field for selecting the targets
target = target.copy_with_fields("labels")
# get the targets corresponding GT for each proposal
matched_targets = target[matched_idxs.clamp(min=0)]
matched_targets.add_field("matched_idxs", matched_idxs)
return matched_targets
def prepare_targets_point(self, proposals, targets):
labels = []
for proposals_per_image, targets_per_image in zip(proposals, targets):
matched_targets = self.match_targets_to_proposals_point(
proposals_per_image, targets_per_image)
matched_idxs = matched_targets.get_field("matched_idxs")
labels_per_image = matched_targets.get_field("labels")
labels_per_image = labels_per_image.to(dtype=torch.int64)
# Label background (below the low threshold)
bg_inds = matched_idxs == Matcher.BELOW_LOW_THRESHOLD
labels_per_image[bg_inds] = 0
# Label ignore proposals (between low and high thresholds)
ignore_inds = matched_idxs == Matcher.BETWEEN_THRESHOLDS
labels_per_image[ignore_inds] = -1 # -1 is ignored by sampler
labels.append(labels_per_image)
return labels
def subsample_point(self, proposals, targets):
labels = self.prepare_targets_point(proposals, targets)
sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
proposals = list(proposals)
# add corresponding label and regression_targets information to the bounding boxes
for labels_per_image, proposals_per_image in zip(labels, proposals):
proposals_per_image.add_field("labels", labels_per_image)
# distributed sampled proposals, that were obtained on all feature maps
# concatenated via the fg_bg_sampler, into individual feature map levels
for img_idx, (pos_inds_img, neg_inds_img) in enumerate(
zip(sampled_pos_inds, sampled_neg_inds)):
img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img,
as_tuple=False).squeeze(1)
proposals_per_image = proposals[img_idx][img_sampled_inds]
proposals[img_idx] = proposals_per_image
return proposals
def match_targets_to_proposals_scribble(self, proposal, target):
match_quality_matrix_async = boxlist_iou_async(
target.get_field('scribble'), proposal)
matched_idxs = self.proposal_matcher(match_quality_matrix_async)
# pseudo_target = proposal[_matched_idxs>0]
# match_quality_matrix = boxlist_iou(pseudo_target, proposal)
# matched_idxs = self.proposal_matcher(match_quality_matrix)
# Fast RCNN only need "labels" field for selecting the targets
target = target.copy_with_fields("labels")
# get the targets corresponding GT for each proposal
matched_targets = target[matched_idxs.clamp(min=0)]
matched_targets.add_field("matched_idxs", matched_idxs)
return matched_targets
def prepare_targets_scribble(self, proposals, targets):
labels = []
for proposals_per_image, targets_per_image in zip(proposals, targets):
matched_targets = self.match_targets_to_proposals_scribble(
proposals_per_image, targets_per_image)
matched_idxs = matched_targets.get_field("matched_idxs")
labels_per_image = matched_targets.get_field("labels")
labels_per_image = labels_per_image.to(dtype=torch.int64)
# Label background (below the low threshold)
bg_inds = matched_idxs == Matcher.BELOW_LOW_THRESHOLD
labels_per_image[bg_inds] = 0
# Label ignore proposals (between low and high thresholds)
ignore_inds = matched_idxs == Matcher.BETWEEN_THRESHOLDS
labels_per_image[ignore_inds] = -1 # -1 is ignored by sampler
labels.append(labels_per_image)
return labels
def subsample_scribble(self, proposals, targets):
labels = self.prepare_targets_scribble(proposals, targets)
sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
proposals = list(proposals)
# add corresponding label and regression_targets information to the bounding boxes
for labels_per_image, proposals_per_image in zip(labels, proposals):
proposals_per_image.add_field("labels", labels_per_image)
# distributed sampled proposals, that were obtained on all feature maps
# concatenated via the fg_bg_sampler, into individual feature map levels
for img_idx, (pos_inds_img, neg_inds_img) in enumerate(
zip(sampled_pos_inds, sampled_neg_inds)):
img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img,
as_tuple=False).squeeze(1)
proposals_per_image = proposals[img_idx][img_sampled_inds]
proposals[img_idx] = proposals_per_image
return proposals
def __call__(self, proposals, targets):
if self.sampler_type == 'box':
return self.subsample_box(proposals, targets)
elif self.sampler_type == "point":
return self.subsample_point(proposals, targets)
elif self.sampler_type == "scribble":
return self.subsample_scribble(proposals, targets)
else:
raise ValueError
def __init__(self, p, iou=0.2):
self.portion = p
self.iou_th = iou
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
class mist_layer(object):
def __init__(self, p, iou=0.2):
self.portion = p
self.iou_th = iou
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
@torch.no_grad()
def __call__(self,
proposals,
source_score,
labels,
device,
return_targets=False):
num_rois = len(proposals)
k = int(num_rois * self.portion)
num_gt_cls = labels[1:].sum()
if num_gt_cls != 0 and num_rois != 0:
cls_prob = source_score[:, 1:]
gt_cls_inds = labels[1:].nonzero(as_tuple=False)[:, 0]
sorted_scores, max_inds = cls_prob[:, gt_cls_inds].sort(
dim=0, descending=True)
sorted_scores = sorted_scores[:k]
max_inds = max_inds[:k]
_boxes = proposals.bbox[max_inds.t().contiguous().view(-1)].view(
num_gt_cls.int(), -1, 4)
_boxes = BatchBoxList(_boxes, proposals.size, mode=proposals.mode)
ious = batch_boxlist_iou(_boxes, _boxes)
k_ind = torch.zeros(num_gt_cls.int(),
k,
dtype=torch.bool,
device=device)
k_ind[:, 0] = 1 # always take the one with max score
for ii in range(1, k):
max_iou, _ = torch.max(ious[:, ii:ii + 1, :ii], dim=2)
k_ind[:, ii] = (max_iou < self.iou_th).byte().squeeze(-1)
gt_boxes = _boxes.bbox[k_ind]
gt_cls_id = gt_cls_inds + 1
temp_cls = torch.ones(
(_boxes.bbox.shape[:2]), device=device) * gt_cls_id.view(
-1, 1).float()
gt_classes = temp_cls[k_ind].view(-1, 1).long()
gt_scores = sorted_scores.t().contiguous()[k_ind].view(-1, 1)
if gt_boxes.shape[0] != 0:
gt_boxes = BoxList(gt_boxes,
proposals.size,
mode=proposals.mode)
overlaps = boxlist_iou(proposals, gt_boxes)
# TODO: pytorch and numpy argmax perform differently
# max_overlaps, gt_assignment = overlaps.max(dim=1)
max_overlaps = torch.tensor(overlaps.cpu().numpy().max(axis=1),
device=device)
gt_assignment = torch.tensor(
overlaps.cpu().numpy().argmax(axis=1), device=device)
pseudo_labels = gt_classes[gt_assignment, 0]
loss_weights = gt_scores[gt_assignment, 0]
# fg_inds = max_overlaps.ge(cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD).nonzero(as_tuple=False)[:,0]
# Select background RoIs as those with <= FG_IOU_THRESHOLD
bg_inds = max_overlaps.lt(
cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD).nonzero(
as_tuple=False)[:, 0]
pseudo_labels[bg_inds] = 0
# compute regression targets
if return_targets:
matched_targets = gt_boxes[gt_assignment]
regression_targets = self.box_coder.encode(
matched_targets.bbox, proposals.bbox)
return pseudo_labels, loss_weights, regression_targets
return pseudo_labels, loss_weights
# corner case
pseudo_labels = torch.zeros(num_rois, dtype=torch.long, device=device)
loss_weights = torch.zeros(num_rois, dtype=torch.float, device=device)
if return_targets:
regression_targets = torch.zeros(num_rois,
4,
dtype=torch.float,
device=device)
return pseudo_labels, loss_weights, regression_targets
return pseudo_labels, loss_weights
def __init__(self):
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
class oicr_layer(object):
""" OICR. Tang et al. 2017 (https://arxiv.org/abs/1704.00138) """
def __init__(self):
bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
self.box_coder = BoxCoder(weights=bbox_reg_weights)
@torch.no_grad()
def __call__(self,
proposals,
source_score,
labels,
device,
return_targets=False):
gt_boxes = torch.zeros((0, 4), dtype=torch.float, device=device)
gt_classes = torch.zeros((0, 1), dtype=torch.long, device=device)
gt_scores = torch.zeros((0, 1), dtype=torch.float, device=device)
# not using the background class
_prob = source_score[:, 1:].clone()
_labels = labels[1:]
#positive_classes = _labels.eq(1).nonzero(as_tuple=False)[:, 0]
positive_classes = torch.arange(
_labels.shape[0])[_labels == 1].to(device)
for c in positive_classes:
cls_prob = _prob[:, c]
max_index = torch.argmax(cls_prob)
gt_boxes = torch.cat(
(gt_boxes, proposals.bbox[max_index].view(1, -1)), dim=0)
gt_classes = torch.cat((gt_classes, c.add(1).view(1, 1)), dim=0)
gt_scores = torch.cat((gt_scores, cls_prob[max_index].view(1, 1)),
dim=0)
_prob[max_index].fill_(0)
#if return_targets == True:
# gt_boxes = BoxList(gt_boxes, proposals.size, mode=proposals.mode)
# gt_boxes.add_field('labels', gt_classes[:, 0].float())
# # gt_boxes.add_field('difficult', bb)
# return gt_boxes
if gt_boxes.shape[0] == 0:
num_rois = len(source_score)
pseudo_labels = torch.zeros(num_rois,
dtype=torch.long,
device=device)
loss_weights = torch.zeros(num_rois,
dtype=torch.float,
device=device)
else:
gt_boxes = BoxList(gt_boxes, proposals.size, mode=proposals.mode)
overlaps = boxlist_iou(proposals, gt_boxes)
max_overlaps, gt_assignment = overlaps.max(dim=1)
pseudo_labels = gt_classes[gt_assignment, 0]
loss_weights = gt_scores[gt_assignment, 0]
# Select background RoIs as those with <= FG_IOU_THRESHOLD
bg_inds = max_overlaps.le(
cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD).nonzero(
as_tuple=False)[:, 0]
pseudo_labels[bg_inds] = 0
if return_targets == True:
matched_targets = gt_boxes[gt_assignment]
regression_targets = self.box_coder.encode(
matched_targets.bbox, proposals.bbox)
return pseudo_labels, loss_weights, regression_targets
# PCL_TRICK:
# ignore_thres = 0.1
# ignore_inds = max_overlaps.le(ignore_thres).nonzero(as_tuple=False)[:,0]
# loss_weights[ignore_inds] = 0
return pseudo_labels, loss_weights