205 lines
7.4 KiB
Python
205 lines
7.4 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Part of the implementation is borrowed and modified from internet,
|
|
# publicly available at https://github.com/CAIC-AD/YOLOPv2
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torchvision.ops import nms
|
|
|
|
|
|
def _make_grid(nx=20, ny=20):
|
|
yv, xv = torch.meshgrid(
|
|
[torch.arange(ny), torch.arange(nx)], indexing='ij')
|
|
return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
|
|
|
|
|
|
def split_for_trace_model(pred=None, anchor_grid=None):
|
|
z = []
|
|
st = [8, 16, 32]
|
|
for i in range(3):
|
|
bs, _, ny, nx = pred[i].shape
|
|
pred[i] = pred[i].view(bs, 3, 85, ny, nx).permute(0, 1, 3, 4,
|
|
2).contiguous()
|
|
y = pred[i].sigmoid()
|
|
gr = _make_grid(nx, ny).to(pred[i].device)
|
|
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + gr) * st[i] # xy
|
|
y[..., 2:4] = (y[..., 2:4] * 2)**2 * anchor_grid[i] # wh
|
|
z.append(y.view(bs, -1, 85))
|
|
pred = torch.cat(z, 1)
|
|
return pred
|
|
|
|
|
|
def scale_coords(img1_shape, coords, img0_shape=(720, 1280), ratio_pad=None):
|
|
# Rescale coords (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]
|
|
|
|
coords[:, [0, 2]] -= pad[0] # x padding
|
|
coords[:, [1, 3]] -= pad[1] # y padding
|
|
coords[:, :4] /= gain
|
|
clip_coords(coords, img0_shape)
|
|
return coords
|
|
|
|
|
|
def clip_coords(boxes, img_shape):
|
|
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
|
boxes[:, 0].clamp_(0, img_shape[1]) # x1
|
|
boxes[:, 1].clamp_(0, img_shape[0]) # y1
|
|
boxes[:, 2].clamp_(0, img_shape[1]) # x2
|
|
boxes[:, 3].clamp_(0, img_shape[0]) # y2
|
|
|
|
|
|
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 non_max_suppression(prediction,
|
|
conf_thres=0.3,
|
|
iou_thres=0.45,
|
|
classes=None,
|
|
agnostic=False,
|
|
multi_label=False,
|
|
labels=()):
|
|
"""Runs Non-Maximum Suppression (NMS) on inference results
|
|
|
|
Returns:
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
"""
|
|
|
|
nc = prediction.shape[2] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Settings
|
|
max_wh = 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
|
time_limit = 10.0 # 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), device=prediction.device)] * prediction.shape[0]
|
|
for xi, x in enumerate(prediction): # image index, image inference
|
|
# Apply constraints
|
|
x = x[xc[xi]] # confidence
|
|
|
|
# Cat apriori labels if autolabelling
|
|
if labels and len(labels[xi]):
|
|
lbs = labels[xi]
|
|
v = torch.zeros((len(lbs), nc + 5), device=x.device)
|
|
v[:, :4] = lbs[:, 1:5] # box
|
|
v[:, 4] = 1.0 # conf
|
|
v[range(len(lbs)), lbs[:, 0].long() + 5] = 1.0 # cls
|
|
x = torch.cat((x, v), 0)
|
|
|
|
# If none remain process next image
|
|
if not x.shape[0]:
|
|
continue
|
|
|
|
# Compute conf
|
|
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
|
|
|
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
|
box = xywh2xyxy(x[:, :4])
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
if multi_label:
|
|
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
|
else: # best class only
|
|
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float()),
|
|
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)]
|
|
|
|
# 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
|
|
|
|
# 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 = nms(boxes, scores, iou_thres) # NMS
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
if merge and (1 < n and n < 3E3):
|
|
# 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 (time.time() - t) > time_limit:
|
|
print(f'WARNING: NMS time limit {time_limit}s exceeded')
|
|
break # time limit exceeded
|
|
|
|
return output
|
|
|
|
|
|
def box_iou(box1, box2):
|
|
# 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.
|
|
Args:
|
|
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
|
|
"""
|
|
|
|
def box_area(box):
|
|
# box = 4xn
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
area1 = box_area(box1.T)
|
|
area2 = box_area(box2.T)
|
|
|
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:])
|
|
- torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
return inter / (area1[:, None] + area2 - inter
|
|
) # iou = inter / (area1 + area2 - inter)
|
|
|
|
|
|
def driving_area_mask(da_predict=None, out_shape=(720, 1280)):
|
|
da_seg_mask = torch.nn.functional.interpolate(
|
|
da_predict, size=out_shape, mode='bilinear')
|
|
_, da_seg_mask = torch.max(da_seg_mask, 1)
|
|
da_seg_mask = da_seg_mask.int().squeeze().cpu().numpy()
|
|
return da_seg_mask
|
|
|
|
|
|
def lane_line_mask(ll_predict=None, out_shape=(720, 1280)):
|
|
ll_seg_mask = torch.nn.functional.interpolate(
|
|
ll_predict, size=out_shape, mode='bilinear')
|
|
ll_seg_mask = torch.round(ll_seg_mask).squeeze(1)
|
|
ll_seg_mask = ll_seg_mask.int().squeeze().cpu().numpy()
|
|
return ll_seg_mask
|