Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import math | |
| def bbox_bep(box1, box2, xywh=False, eps=1e-7, bep1 = True): | |
| """ | |
| Calculates bottom edge proximity between two boxes | |
| Input shapes are box1(1,4) to box2(n,4) | |
| Implementation of bep2 from | |
| Are object detection assessment criteria ready for maritime computer vision? | |
| """ | |
| # Get the coordinates of bounding boxes | |
| if xywh: # transform from xywh to xyxy | |
| (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1) | |
| w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2 | |
| b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_ | |
| b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_ | |
| else: # x1, y1, x2, y2 = box1 | |
| b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1) | |
| b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1) | |
| w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps) | |
| w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps) | |
| # Bottom edge distance (absolute value) | |
| # xb = torch.abs(b2_x2 - b1_x1) | |
| xb = torch.min(b2_x2-b1_x1, b1_x2-b2_x1) | |
| xa = w2 - xb | |
| xc = w1 - xb | |
| ybe = torch.abs(b2_y2 - b1_y2) | |
| X2 = xb/(xb+xa) | |
| Y2 = 1-ybe/h2 | |
| X1 = xb/(xb+xa+xc+eps) | |
| Y1 = 1-ybe/(torch.max(h2,h1)+eps) | |
| bep = X1*Y1 if bep1 else X2*Y2 | |
| return bep | |
| def bbox_iou(box1, box2, xywh=False, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): | |
| """ | |
| Calculates IoU, GIoU, DIoU, or CIoU between two boxes, supporting xywh/xyxy formats. | |
| Input shapes are box1(1,4) to box2(n,4). | |
| """ | |
| # Get the coordinates of bounding boxes | |
| if xywh: # transform from xywh to xyxy | |
| (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1) | |
| w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2 | |
| b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_ | |
| b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_ | |
| else: # x1, y1, x2, y2 = box1 | |
| b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1) | |
| b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1) | |
| w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps) | |
| w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps) | |
| # Intersection area | |
| inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * ( | |
| b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1) | |
| ).clamp(0) | |
| # Union Area | |
| union = w1 * h1 + w2 * h2 - inter + eps | |
| # IoU | |
| iou = inter / union | |
| if CIoU or DIoU or GIoU: | |
| cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1) # convex (smallest enclosing box) width | |
| ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1) # convex height | |
| if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 | |
| c2 = cw**2 + ch**2 + eps # convex diagonal squared | |
| rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center dist ** 2 | |
| if CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 | |
| v = (4 / math.pi**2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2) | |
| with torch.no_grad(): | |
| alpha = v / (v - iou + (1 + eps)) | |
| return iou - (rho2 / c2 + v * alpha) # CIoU | |
| return iou - rho2 / c2 # DIoU | |
| c_area = cw * ch + eps # convex area | |
| return iou - (c_area - union) / c_area # GIoU https://arxiv.org/pdf/1902.09630.pdf | |
| return iou # IoU | |