| 3 | import math |
| 4 | |
| 5 | class BoxCoder(object): |
| 6 | |
| 7 | def __init__(self, opt): |
| 8 | self.cfg = opt |
| 9 | |
| 10 | def encode(self, gt_boxes, anchors): |
| 11 | if False: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'POINT': |
| 12 | TO_REMOVE = 1 # TODO remove |
| 13 | anchors_w = anchors[:, 2] - anchors[:, 0] + TO_REMOVE |
| 14 | anchors_h = anchors[:, 3] - anchors[:, 1] + TO_REMOVE |
| 15 | anchors_cx = (anchors[:, 2] + anchors[:, 0]) / 2 |
| 16 | anchors_cy = (anchors[:, 3] + anchors[:, 1]) / 2 |
| 17 | |
| 18 | w = self.cfg.MODEL.ATSS.ANCHOR_SIZES[0] / self.cfg.MODEL.ATSS.ANCHOR_STRIDES[0] |
| 19 | l = w * (anchors_cx - gt_boxes[:, 0]) / anchors_w |
| 20 | t = w * (anchors_cy - gt_boxes[:, 1]) / anchors_h |
| 21 | r = w * (gt_boxes[:, 2] - anchors_cx) / anchors_w |
| 22 | b = w * (gt_boxes[:, 3] - anchors_cy) / anchors_h |
| 23 | targets = torch.stack([l, t, r, b], dim=1) |
| 24 | elif True: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'BOX': |
| 25 | TO_REMOVE = 1 # TODO remove |
| 26 | ex_length = anchors[:, 1] - anchors[:, 0] + TO_REMOVE |
| 27 | ex_center = (anchors[:, 1] + anchors[:, 0]) / 2 |
| 28 | |
| 29 | gt_length = gt_boxes[:, 1] - gt_boxes[:, 0] + TO_REMOVE |
| 30 | gt_center = (gt_boxes[:, 1] + gt_boxes[:, 0]) / 2 |
| 31 | |
| 32 | wx, ww = (10., 5.) |
| 33 | targets_dx = wx * (gt_center - ex_center) / ex_length |
| 34 | targets_dw = ww * torch.log(gt_length / ex_length) |
| 35 | targets = torch.stack((targets_dx, targets_dw), dim=1) |
| 36 | |
| 37 | return targets |
| 38 | |
| 39 | def decode(self, preds, anchors): |
| 40 | if False: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'POINT': |
| 41 | TO_REMOVE = 1 # TODO remove |
| 42 | anchors_w = anchors[:, 2] - anchors[:, 0] + TO_REMOVE |
| 43 | anchors_h = anchors[:, 3] - anchors[:, 1] + TO_REMOVE |
| 44 | anchors_cx = (anchors[:, 2] + anchors[:, 0]) / 2 |
| 45 | anchors_cy = (anchors[:, 3] + anchors[:, 1]) / 2 |
| 46 | |
| 47 | w = self.cfg.MODEL.ATSS.ANCHOR_SIZES[0] / self.cfg.MODEL.ATSS.ANCHOR_STRIDES[0] |
| 48 | x1 = anchors_cx - preds[:, 0] / w * anchors_w |
| 49 | y1 = anchors_cy - preds[:, 1] / w * anchors_h |
| 50 | x2 = anchors_cx + preds[:, 2] / w * anchors_w |
| 51 | y2 = anchors_cy + preds[:, 3] / w * anchors_h |
| 52 | pred_boxes = torch.stack([x1, y1, x2, y2], dim=1) |
| 53 | elif True: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'BOX': |
| 54 | anchors = anchors.to(preds.dtype) |
| 55 | |
| 56 | TO_REMOVE = 1 # TODO remove |
| 57 | ex_length = anchors[:, 1] - anchors[:, 0] + TO_REMOVE |
| 58 | ex_center = (anchors[:, 1] + anchors[:, 0]) / 2 |
| 59 | |
| 60 | wx, ww = (10, 5.) |
| 61 | dx = preds[:, 0] / wx |
| 62 | dw = preds[:, 1] / ww |