Fix class-agnostic bbox regression

Reviewed By: ir413

Differential Revision: D8083267

fbshipit-source-id: 96c17d79d117bf434a5b879a8bca4c88ea6b68e2

detectron/datasets/roidb.py

@@ -139,10 +139,10 @@ def filter_for_training(roidb):
 def add_bbox_regression_targets(roidb):
     """Add information needed to train bounding-box regressors."""
     for entry in roidb:
-        entry['bbox_targets'] = _compute_targets(entry)
+        entry['bbox_targets'] = compute_bbox_regression_targets(entry)
 
 
-def _compute_targets(entry):
+def compute_bbox_regression_targets(entry):
     """Compute bounding-box regression targets for an image."""
     # Indices of ground-truth ROIs
     rois = entry['boxes']

detectron/modeling/fast_rcnn_heads.py

@@ -45,6 +45,7 @@ import detectron.utils.blob as blob_utils
 
 def add_fast_rcnn_outputs(model, blob_in, dim):
     """Add RoI classification and bounding box regression output ops."""
+    # Box classification layer
     model.FC(
         blob_in,
         'cls_score',
@@ -57,11 +58,15 @@ def add_fast_rcnn_outputs(model, blob_in, dim):
         # Only add softmax when testing; during training the softmax is combined
         # with the label cross entropy loss for numerical stability
         model.Softmax('cls_score', 'cls_prob', engine='CUDNN')
+    # Box regression layer
+    num_bbox_reg_classes = (
+        2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else model.num_classes
+    )
     model.FC(
         blob_in,
         'bbox_pred',
         dim,
-        model.num_classes * 4,
+        num_bbox_reg_classes * 4,
         weight_init=gauss_fill(0.001),
         bias_init=const_fill(0.0)
     )

detectron/modeling/rfcn_heads.py

@@ -55,7 +55,7 @@ def add_rfcn_outputs(model, blob_in, dim_in, dim_reduce, spatial_scale):
         weight_init=gauss_fill(0.01),
         bias_init=const_fill(0.0)
     )
-    # # Bounding-box regression conv
+    # Bounding-box regression conv
     num_bbox_reg_classes = (
         2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else model.num_classes
     )

detectron/ops/collect_and_distribute_fpn_rpn_proposals.py

@@ -22,6 +22,7 @@ import numpy as np
 
 from detectron.core.config import cfg
 from detectron.datasets import json_dataset
+from detectron.datasets import roidb as roidb_utils
 import detectron.modeling.FPN as fpn
 import detectron.roi_data.fast_rcnn as fast_rcnn_roi_data
 import detectron.utils.blob as blob_utils
@@ -53,6 +54,7 @@ class CollectAndDistributeFpnRpnProposalsOp(object):
             # This choice should be investigated in the future (it likely does
             # not matter).
             json_dataset.add_proposals(roidb, rois, im_scales, crowd_thresh=0)
+            roidb_utils.add_bbox_regression_targets(roidb)
             # Compute training labels for the RPN proposals; also handles
             # distributing the proposals over FPN levels
             output_blob_names = fast_rcnn_roi_data.get_fast_rcnn_blob_names()

detectron/ops/generate_proposal_labels.py

@@ -21,6 +21,7 @@ from __future__ import unicode_literals
 import logging
 
 from detectron.datasets import json_dataset
+from detectron.datasets import roidb as roidb_utils
 from detectron.utils import blob as blob_utils
 import detectron.roi_data.fast_rcnn as fast_rcnn_roi_data
 
@@ -46,6 +47,7 @@ class GenerateProposalLabelsOp(object):
         # This choice should be investigated in the future (it likely does
         # not matter).
         json_dataset.add_proposals(roidb, rois, im_scales, crowd_thresh=0)
+        roidb_utils.add_bbox_regression_targets(roidb)
         blobs = {k: [] for k in output_blob_names}
         fast_rcnn_roi_data.add_fast_rcnn_blobs(blobs, im_scales, roidb)
         for i, k in enumerate(output_blob_names):

detectron/roi_data/fast_rcnn.py

@@ -170,19 +170,9 @@ def _sample_rois(roidb, im_scale, batch_idx):
     sampled_labels[fg_rois_per_this_image:] = 0  # Label bg RoIs with class 0
     sampled_boxes = roidb['boxes'][keep_inds]
 
-    if 'bbox_targets' not in roidb:
-        gt_inds = np.where(roidb['gt_classes'] > 0)[0]
-        gt_boxes = roidb['boxes'][gt_inds, :]
-        gt_assignments = gt_inds[roidb['box_to_gt_ind_map'][keep_inds]]
-        bbox_targets = _compute_targets(
-            sampled_boxes, gt_boxes[gt_assignments, :], sampled_labels
-        )
-        bbox_targets, bbox_inside_weights = _expand_bbox_targets(bbox_targets)
-    else:
-        bbox_targets, bbox_inside_weights = _expand_bbox_targets(
-            roidb['bbox_targets'][keep_inds, :]
-        )
-
+    bbox_targets, bbox_inside_weights = _expand_bbox_targets(
+        roidb['bbox_targets'][keep_inds, :]
+    )
     bbox_outside_weights = np.array(
         bbox_inside_weights > 0, dtype=bbox_inside_weights.dtype
     )
@@ -216,21 +206,6 @@ def _sample_rois(roidb, im_scale, batch_idx):
     return blob_dict
 
 
-def _compute_targets(ex_rois, gt_rois, labels):
-    """Compute bounding-box regression targets for an image."""
-
-    assert ex_rois.shape[0] == gt_rois.shape[0]
-    assert ex_rois.shape[1] == 4
-    assert gt_rois.shape[1] == 4
-
-    targets = box_utils.bbox_transform_inv(
-        ex_rois, gt_rois, cfg.MODEL.BBOX_REG_WEIGHTS
-    )
-    return np.hstack((labels[:, np.newaxis], targets)).astype(
-        np.float32, copy=False
-    )
-
-
 def _expand_bbox_targets(bbox_target_data):
     """Bounding-box regression targets are stored in a compact form in the
     roidb.