244 lines
10 KiB
Python
244 lines
10 KiB
Python
import cv2
|
|
import torch
|
|
import torchvision
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
import time
|
|
import logging
|
|
import logging.config
|
|
TORCH_1_10 = False
|
|
LOGGING_NAME = "yolov9"
|
|
LOGGER = logging.getLogger(LOGGING_NAME)
|
|
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
|
# Resize and pad image while meeting stride-multiple constraints
|
|
shape = im.shape[:2] # current shape [height, width]
|
|
if isinstance(new_shape, int):
|
|
new_shape = (new_shape, new_shape)
|
|
|
|
# Scale ratio (new / old)
|
|
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
|
if not scaleup: # only scale down, do not scale up (for better val mAP)
|
|
r = min(r, 1.0)
|
|
|
|
# Compute padding
|
|
ratio = r, r # width, height ratios
|
|
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
|
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
|
if auto: # minimum rectangle
|
|
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
|
elif scaleFill: # stretch
|
|
dw, dh = 0.0, 0.0
|
|
new_unpad = (new_shape[1], new_shape[0])
|
|
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
|
|
|
|
dw /= 2 # divide padding into 2 sides
|
|
dh /= 2
|
|
|
|
if shape[::-1] != new_unpad: # resize
|
|
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
|
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
|
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
|
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
|
return im, ratio, (dw, dh)
|
|
|
|
|
|
def make_anchors(feats, strides, grid_cell_offset=0.5):
|
|
"""Generate anchors from features."""
|
|
anchor_points, stride_tensor = [], []
|
|
assert feats is not None
|
|
dtype, device = feats[0].dtype, feats[0].device
|
|
for i, stride in enumerate(strides):
|
|
_, _, h, w = feats[i].shape
|
|
sx = torch.arange(end=w, device=device, dtype=dtype) + grid_cell_offset # shift x
|
|
sy = torch.arange(end=h, device=device, dtype=dtype) + grid_cell_offset # shift y
|
|
sy, sx = torch.meshgrid(sy, sx, indexing='ij') if TORCH_1_10 else torch.meshgrid(sy, sx)
|
|
anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
|
|
stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
|
|
return torch.cat(anchor_points), torch.cat(stride_tensor)
|
|
|
|
def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
|
|
"""Transform distance(ltrb) to box(xywh or xyxy)."""
|
|
lt, rb = torch.split(distance, 2, dim)
|
|
x1y1 = anchor_points - lt
|
|
x2y2 = anchor_points + rb
|
|
if xywh:
|
|
c_xy = (x1y1 + x2y2) / 2
|
|
wh = x2y2 - x1y1
|
|
return torch.cat((c_xy, wh), dim) # xywh bbox
|
|
return torch.cat((x1y1, x2y2), dim) # xyxy bbox
|
|
|
|
class DFL(nn.Module):
|
|
# DFL module
|
|
def __init__(self, c1=17):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
|
|
self.conv.weight.data[:] = nn.Parameter(torch.arange(c1, dtype=torch.float).view(1, c1, 1, 1)) # / 120.0
|
|
self.c1 = c1
|
|
# self.bn = nn.BatchNorm2d(4)
|
|
|
|
def forward(self, x):
|
|
b, c, a = x.shape # batch, channels, anchors
|
|
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
|
|
# return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
|
|
|
|
def xywh2xyxy(x):
|
|
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
|
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
|
|
y[..., 1] = x[..., 1] - x[..., 3] / 2 # top left y
|
|
y[..., 2] = x[..., 0] + x[..., 2] / 2 # bottom right x
|
|
y[..., 3] = x[..., 1] + x[..., 3] / 2 # bottom right y
|
|
return y
|
|
def box_iou(box1, box2, eps=1e-7):
|
|
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
|
"""
|
|
Return intersection-over-union (Jaccard index) of boxes.
|
|
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
|
Arguments:
|
|
box1 (Tensor[N, 4])
|
|
box2 (Tensor[M, 4])
|
|
Returns:
|
|
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
|
IoU values for every element in boxes1 and boxes2
|
|
"""
|
|
|
|
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
|
(a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
|
|
inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
|
|
|
|
# IoU = inter / (area1 + area2 - inter)
|
|
return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
|
|
def non_max_suppression(
|
|
prediction,
|
|
conf_thres=0.25,
|
|
iou_thres=0.45,
|
|
classes=None,
|
|
agnostic=False,
|
|
multi_label=False,
|
|
labels=(),
|
|
max_det=300,
|
|
nm=0, # number of masks
|
|
):
|
|
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
|
|
|
|
Returns:
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
"""
|
|
|
|
if isinstance(prediction, (list, tuple)): # YOLO model in validation model, output = (inference_out, loss_out)
|
|
prediction = prediction[0] # select only inference output
|
|
|
|
device = prediction.device
|
|
mps = 'mps' in device.type # Apple MPS
|
|
if mps: # MPS not fully supported yet, convert tensors to CPU before NMS
|
|
prediction = prediction.cpu()
|
|
bs = prediction.shape[0] # batch size
|
|
nc = prediction.shape[1] - nm - 4 # number of classes
|
|
mi = 4 + nc # mask start index
|
|
xc = prediction[:, 4:mi].amax(1) > conf_thres # candidates
|
|
|
|
# Checks
|
|
assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
|
|
assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
|
|
|
|
# Settings
|
|
# min_wh = 2 # (pixels) minimum box width and height
|
|
max_wh = 7680 # (pixels) maximum box width and height
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
|
time_limit = 2.5 + 0.05 * bs # seconds to quit after
|
|
redundant = True # require redundant detections
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
merge = False # use merge-NMS
|
|
|
|
t = time.time()
|
|
output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
|
|
for xi, x in enumerate(prediction): # image index, image inference
|
|
# Apply constraints
|
|
# x[((x[:, 2:4] < min_wh) | (x[:, 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
|
x = x.T[xc[xi]] # confidence
|
|
|
|
# Cat apriori labels if autolabelling
|
|
if labels and len(labels[xi]):
|
|
lb = labels[xi]
|
|
v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
|
|
v[:, :4] = lb[:, 1:5] # box
|
|
v[range(len(lb)), lb[:, 0].long() + 4] = 1.0 # cls
|
|
x = torch.cat((x, v), 0)
|
|
|
|
# If none remain process next image
|
|
if not x.shape[0]:
|
|
continue
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
box, cls, mask = x.split((4, nc, nm), 1)
|
|
box = xywh2xyxy(box) # center_x, center_y, width, height) to (x1, y1, x2, y2)
|
|
if multi_label:
|
|
i, j = (cls > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, 4 + j, None], j[:, None].float(), mask[i]), 1)
|
|
else: # best class only
|
|
conf, j = cls.max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
|
|
|
|
# Filter by class
|
|
if classes is not None:
|
|
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
|
|
|
# Apply finite constraint
|
|
# if not torch.isfinite(x).all():
|
|
# x = x[torch.isfinite(x).all(1)]
|
|
|
|
# Check shape
|
|
n = x.shape[0] # number of boxes
|
|
if not n: # no boxes
|
|
continue
|
|
elif n > max_nms: # excess boxes
|
|
x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
|
else:
|
|
x = x[x[:, 4].argsort(descending=True)] # sort by confidence
|
|
|
|
# Batched NMS
|
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
|
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
|
weights = iou * scores[None] # box weights
|
|
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
|
if redundant:
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
|
|
output[xi] = x[i]
|
|
if mps:
|
|
output[xi] = output[xi].to(device)
|
|
if (time.time() - t) > time_limit:
|
|
LOGGER.warning(f'WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded')
|
|
break # time limit exceeded
|
|
|
|
return output
|
|
def clip_boxes(boxes, shape):
|
|
# Clip boxes (xyxy) to image shape (height, width)
|
|
if isinstance(boxes, torch.Tensor): # faster individually
|
|
boxes[:, 0].clamp_(0, shape[1]) # x1
|
|
boxes[:, 1].clamp_(0, shape[0]) # y1
|
|
boxes[:, 2].clamp_(0, shape[1]) # x2
|
|
boxes[:, 3].clamp_(0, shape[0]) # y2
|
|
else: # np.array (faster grouped)
|
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
|
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
|
def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
|
|
# Rescale boxes (xyxy) from img1_shape to img0_shape
|
|
if ratio_pad is None: # calculate from img0_shape
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
|
else:
|
|
gain = ratio_pad[0][0]
|
|
pad = ratio_pad[1]
|
|
|
|
boxes[:, [0, 2]] -= pad[0] # x padding
|
|
boxes[:, [1, 3]] -= pad[1] # y padding
|
|
boxes[:, :4] /= gain
|
|
clip_boxes(boxes, img0_shape)
|
|
return boxes |