def __init__(self, cfg):
self.center_form_priors = PriorBox(cfg)()
self.corner_form_priors = box_utils.center_form_to_corner_form(
self.center_form_priors)
self.center_variance = cfg.MODEL.CENTER_VARIANCE
self.size_variance = cfg.MODEL.SIZE_VARIANCE
self.iou_threshold = cfg.MODEL.THRESHOLD
def build_target_transform(cfg):
if cfg.MODEL.BOX_HEAD.NAME == "SSDRotateBoxHead":
transform = SSDTargetTransform(
DirectionalPriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE,
cfg.MODEL.SIZE_VARIANCE, cfg.MODEL.THRESHOLD)
else:
transform = SSDTargetTransform(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE,
cfg.MODEL.SIZE_VARIANCE, cfg.MODEL.THRESHOLD)
return transform
def _forward_test(self, cls_logits, bbox_pred):
if self.priors is None:
self.priors = PriorBox(self.cfg)().to(bbox_pred.device)
scores = F.softmax(cls_logits, dim=2)
boxes = box_utils.convert_locations_to_boxes(
bbox_pred, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE
)
boxes = box_utils.center_form_to_corner_form(boxes)
detections = (scores, boxes)
detections = self.post_processor(detections)
return detections, {}
def build_target_transform(cfg):
transform = SSDTargetCornerBoxTransform(
PriorCornerBox(cfg)(), cfg.MODEL.COORDINATE_VARIANCE,
cfg.MODEL.THRESHOLD)
return transform
transform = SSDTargetTransform(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE,
cfg.MODEL.THRESHOLD)
return transform