File size: 4,343 Bytes
91ef820 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
import torch
from torchvision.ops.boxes import box_area
import random
import time
def bbox_iou(box1, box2, x1y1x2y2=True):
"""
Returns the IoU of two bounding boxes
"""
if x1y1x2y2:
# Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
else:
# Transform from center and width to exact coordinates
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
# get the coordinates of the intersection rectangle
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
# Intersection area
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, 0) * torch.clamp(inter_rect_y2 - inter_rect_y1, 0)
# Union Area
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
# print(box1, box1.shape)
# print(box2, box2.shape)
return inter_area / (b1_area + b2_area - inter_area + 1e-16)
def xywh2xyxy(x):
x_c, y_c, w, h = x.unbind(-1)
b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
(x_c + 0.5 * w), (y_c + 0.5 * h)]
return torch.stack(b, dim=-1)
def xyxy2xywh(x):
x0, y0, x1, y1 = x.unbind(-1)
b = [(x0 + x1) / 2.0, (y0 + y1) / 2.0,
(x1 - x0), (y1 - y0)]
return torch.stack(b, dim=-1)
def box_iou(boxes1, boxes2):
area1 = box_area(boxes1)
area2 = box_area(boxes2)
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
wh = (rb - lt).clamp(min=0) # [N,M,2]
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
union = area1[:, None] + area2 - inter
iou = inter / union
return iou, union
def generalized_box_iou(boxes1, boxes2):
"""
Generalized IoU from https://giou.stanford.edu/
The boxes should be in [x0, y0, x1, y1] format
Returns a [N, M] pairwise matrix, where N = len(boxes1)
and M = len(boxes2)
"""
# degenerate boxes gives inf / nan results
# so do an early check
assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
iou, union = box_iou(boxes1, boxes2)
lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
wh = (rb - lt).clamp(min=0) # [N,M,2]
area = wh[:, :, 0] * wh[:, :, 1]
return iou - (area - union) / area
def sampleNegBBox(box, CAsampleType, CAsampleNum, category=0, w=640, h=640):
assert CAsampleType in ['random', 'attention', 'crossImage', 'crossBatch']
index = 0
negBox_list = []
# ori_center = [(box[0]+box[2])/2, (box[1]+box[3])/2]
ori_w, ori_h = box[2]-box[0], box[3]-box[1]
flag=0
while index < CAsampleNum:
flag += 1
# print(flag)
if CAsampleType == 'random':
xNeg = torch.randint(1, w, (1,))
yNeg = torch.randint(1, h, (1,))
wNeg = ori_w + random.randint(torch.round(-ori_w * 0.1), torch.round(ori_w * 0.1))
hNeg = ori_h + random.randint(torch.round(-ori_h * 0.1), torch.round(ori_h * 0.1))
elif CAsampleType == 'attention':
pass
negBox = torch.zeros([4])
negBox[0], negBox[1], negBox[2], negBox[3] = xNeg - 0.5 * wNeg, yNeg - 0.5 * hNeg, xNeg + 0.5 * wNeg, yNeg + 0.5 * hNeg
negBox = torch.round(negBox)
# 加入越界条件筛选 invalid bbox
if negBox[0] < 0 or negBox[1] < 0 or negBox[2] >= w or negBox[3] >= h:
continue
# 加入box冲突条件筛选 invalid bbox
iou, union = box_iou(box.unsqueeze(0), negBox.unsqueeze(0))
if iou > 0.25 and flag < 300:
continue
negBox_list.append(negBox)
index += 1
return negBox_list
|