Encodes boxes with respect to anchors. Args: boxes: A [..., 4] float tensor with boxes to encode. Boxes must be of the form [ymin, xmin, ymax, xmax]. anchors: A [..., 4] float tensor containing anchors. Anchors must be broadcastable to
(self, *, boxes: Tensor, anchors: Tensor)
| 54 | clip_boxes: BoxClipMethod = BoxClipMethod.MaxHW |
| 55 | |
| 56 | def encode(self, *, boxes: Tensor, anchors: Tensor) -> Tensor: |
| 57 | """Encodes boxes with respect to anchors. |
| 58 | |
| 59 | Args: |
| 60 | boxes: A [..., 4] float tensor with boxes to encode. Boxes must be of the form |
| 61 | [ymin, xmin, ymax, xmax]. |
| 62 | anchors: A [..., 4] float tensor containing anchors. Anchors must be broadcastable to |
| 63 | `boxes` and of the form [ymin, xmin, ymax, xmax]. |
| 64 | |
| 65 | Returns: |
| 66 | A [..., 4] float tensor containing encoded boxes. |
| 67 | """ |
| 68 | cfg = self.config |
| 69 | boxes = boxes.astype(anchors.dtype) |
| 70 | ymin = boxes[..., 0:1] |
| 71 | xmin = boxes[..., 1:2] |
| 72 | ymax = boxes[..., 2:3] |
| 73 | xmax = boxes[..., 3:4] |
| 74 | box_h = ymax - ymin + cfg.eps |
| 75 | box_w = xmax - xmin + cfg.eps |
| 76 | |
| 77 | box_h = jnp.maximum(box_h, cfg.box_wh_min_value) |
| 78 | box_w = jnp.maximum(box_w, cfg.box_wh_min_value) |
| 79 | |
| 80 | box_yc = ymin + 0.5 * box_h |
| 81 | box_xc = xmin + 0.5 * box_w |
| 82 | |
| 83 | anchor_ymin = anchors[..., 0:1] |
| 84 | anchor_xmin = anchors[..., 1:2] |
| 85 | anchor_ymax = anchors[..., 2:3] |
| 86 | anchor_xmax = anchors[..., 3:4] |
| 87 | anchor_h = anchor_ymax - anchor_ymin + cfg.eps |
| 88 | anchor_w = anchor_xmax - anchor_xmin + cfg.eps |
| 89 | anchor_yc = anchor_ymin + 0.5 * anchor_h |
| 90 | anchor_xc = anchor_xmin + 0.5 * anchor_w |
| 91 | |
| 92 | encoded_dy = (box_yc - anchor_yc) / anchor_h |
| 93 | encoded_dx = (box_xc - anchor_xc) / anchor_w |
| 94 | encoded_dh = jnp.log(box_h / anchor_h) |
| 95 | encoded_dw = jnp.log(box_w / anchor_w) |
| 96 | |
| 97 | encoded_dy *= cfg.weights[0] |
| 98 | encoded_dx *= cfg.weights[1] |
| 99 | encoded_dh *= cfg.weights[2] |
| 100 | encoded_dw *= cfg.weights[3] |
| 101 | |
| 102 | encoded_boxes = jnp.concatenate([encoded_dy, encoded_dx, encoded_dh, encoded_dw], axis=-1) |
| 103 | return encoded_boxes |
| 104 | |
| 105 | def decode(self, *, encoded_boxes: Tensor, anchors: Tensor) -> Tensor: |
| 106 | """Decodes encoded boxes with respect to anchors. |