Replace copy_ op in Box2BoxTransform.apply_deltas

facebookresearch · Mar 2, 2021 · 998bc55 · 998bc55
1 parent e40ddbb
commit 998bc55
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/detectron2/modeling/box_regression.py b/detectron2/modeling/box_regression.py
@@ -107,12 +107,12 @@ def apply_deltas(self, deltas, boxes):
         pred_w = torch.exp(dw) * widths[:, None]
         pred_h = torch.exp(dh) * heights[:, None]
 
-        pred_boxes = torch.zeros_like(deltas)
-        pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w  # x1
-        pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h  # y1
-        pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w  # x2
-        pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h  # y2
-        return pred_boxes
+        x1 = pred_ctr_x - 0.5 * pred_w
+        y1 = pred_ctr_y - 0.5 * pred_h
+        x2 = pred_ctr_x + 0.5 * pred_w
+        y2 = pred_ctr_y + 0.5 * pred_h
+        pred_boxes = torch.stack((x1, y1, x2, y2), dim=-1)
+        return pred_boxes.reshape(deltas.shape)
 
 
 @torch.jit.script