We resize both tensors to [A,B,2] without new malloc: [A,2] -> [A,1,2] -> [A,B,2] [B,2] -> [1,B,2] -> [A,B,2] Then we compute the area of intersect between box_a and box_b. Args: box_a: (tensor) bounding boxes, Shape: [A,4]. box_b: (tensor) bounding boxes, Shape: [B,4].
(box_a, box_b)
| 27 | |
| 28 | |
| 29 | def intersect(box_a, box_b): |
| 30 | """ We resize both tensors to [A,B,2] without new malloc: |
| 31 | [A,2] -> [A,1,2] -> [A,B,2] |
| 32 | [B,2] -> [1,B,2] -> [A,B,2] |
| 33 | Then we compute the area of intersect between box_a and box_b. |
| 34 | Args: |
| 35 | box_a: (tensor) bounding boxes, Shape: [A,4]. |
| 36 | box_b: (tensor) bounding boxes, Shape: [B,4]. |
| 37 | Return: |
| 38 | (tensor) intersection area, Shape: [A,B]. |
| 39 | """ |
| 40 | A = box_a.size(0) |
| 41 | B = box_b.size(0) |
| 42 | max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), |
| 43 | box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) |
| 44 | min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), |
| 45 | box_b[:, :2].unsqueeze(0).expand(A, B, 2)) |
| 46 | inter = torch.clamp((max_xy - min_xy), min=0) |
| 47 | return inter[:, :, 0] * inter[:, :, 1] |
| 48 | |
| 49 | |
| 50 | def jaccard(box_a, box_b): |