334 lines
13 KiB
Python
334 lines
13 KiB
Python
# coding: utf-8
|
|
import torch
|
|
import torchvision
|
|
import time
|
|
import numpy as np
|
|
import sys
|
|
np.set_printoptions(threshold=sys.maxsize)
|
|
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.
|
|
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
|
|
"""
|
|
|
|
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(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
|
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 xywh2xyxy(x):
|
|
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
|
y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(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.1, iou_thres=0.6, top_k_num=3000, merge=False, classes=None, agnostic=False):
|
|
"""Performs Non-Maximum Suppression (NMS) on inference results
|
|
|
|
Returns:
|
|
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
|
"""
|
|
# print('conf_thres',conf_thres)
|
|
if prediction.dtype is torch.float16:
|
|
prediction = prediction.float() # to FP32
|
|
|
|
nc = prediction[0].shape[1] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Settings
|
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
time_limit = 10.0 # seconds to quit after
|
|
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
|
|
t = time.time()
|
|
output = [None] * prediction.shape[0]
|
|
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[xc[xi]] # confidence
|
|
|
|
# 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().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]
|
|
|
|
|
|
|
|
# If none remain process next image
|
|
n = x.shape[0] # number of boxes
|
|
if not n:
|
|
continue
|
|
|
|
# Sort by confidence
|
|
# x = x[x[:, 4].argsort(descending=True)]
|
|
|
|
# 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
|
|
# # Sort by confidence
|
|
ind_Sort_by_confidence = x[:, 4].argsort(descending=True)
|
|
boxes = boxes[ind_Sort_by_confidence][:top_k_num] #
|
|
scores = scores[ind_Sort_by_confidence][:top_k_num] #
|
|
x = x[ind_Sort_by_confidence][:top_k_num] #
|
|
# cross classes nms
|
|
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
|
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 (time.time() - t) > time_limit:
|
|
break # time limit exceeded
|
|
|
|
return output
|
|
|
|
def non_max_suppression_kneron(prediction, conf_thres=0.1, iou_thres=0.6, top_k_num=3000, merge=False, classes=None, agnostic=False):
|
|
"""Performs Non-Maximum Suppression (NMS) on inference results
|
|
|
|
Returns:
|
|
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
|
"""
|
|
if prediction.dtype is torch.float16:
|
|
prediction = prediction.float() # to FP32
|
|
|
|
nc = prediction[0].shape[1] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Settings
|
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
time_limit = 10.0 # seconds to quit after
|
|
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
|
|
t = time.time()
|
|
output = [None] * prediction.shape[0]
|
|
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[xc[xi]] # confidence
|
|
|
|
# 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().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]
|
|
|
|
|
|
|
|
# If none remain process next image
|
|
n = x.shape[0] # number of boxes
|
|
if not n:
|
|
continue
|
|
|
|
# Sort by confidence
|
|
# x = x[x[:, 4].argsort(descending=True)]
|
|
|
|
# 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
|
|
# Sort by confidence
|
|
ind_Sort_by_confidence = x[:, 4].argsort(descending=True)
|
|
boxes = boxes[ind_Sort_by_confidence][:top_k_num] #
|
|
scores = scores[ind_Sort_by_confidence][:top_k_num] #
|
|
x = x[ind_Sort_by_confidence][:top_k_num] #
|
|
# cross classes nms
|
|
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
|
|
|
|
output[xi] = x[i]
|
|
if (time.time() - t) > time_limit:
|
|
break # time limit exceeded
|
|
|
|
return output
|
|
|
|
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 scale_coords_ori(img1_shape, coords, img0_shape, 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 scale_coords(img1_shape, coords, img0_shape, 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 scale_coords_test(img1_shape, coords, img0_shape, ratio_pad=None):
|
|
|
|
coords[:, 0] /= (img1_shape[1] / img0_shape[1])
|
|
coords[:, 2] /= (img1_shape[1] / img0_shape[1])
|
|
coords[:, 1] /= (img1_shape[0] / img0_shape[0])
|
|
coords[:, 3] /= (img1_shape[0] / img0_shape[0])
|
|
clip_coords(coords, img0_shape)
|
|
return coords
|
|
|
|
def classes_mapping(det, num_classes):
|
|
det[:, 5] = det[:, 5] + 1.0
|
|
|
|
def Yolov5_postprocess(pred, img_shape, im0_shape, conf_thres, iou_thres, top_k_num, num_classes, vanish_point, e2e_coco) :
|
|
classes, agnostic_nms = None, False#
|
|
img_h = im0_shape[0]
|
|
vanish_y2 = vanish_point * float(img_h)
|
|
# Apply NMS
|
|
pred = non_max_suppression(pred, conf_thres, iou_thres, top_k_num, classes=classes, agnostic=agnostic_nms)
|
|
#return pred
|
|
dets = []
|
|
for i, det in enumerate(pred): # detections per image
|
|
gn = torch.tensor(im0_shape)[[1, 0, 1, 0]] # normalization gain whwh
|
|
if det is not None and len(det):
|
|
# Rescale boxes from img_size to im0 size
|
|
det[:, :4] = scale_coords(img_shape[2:], det[:, :4], im0_shape).round()
|
|
det = det[det[:,3]>=vanish_y2]
|
|
# (x1,y1,x2,y2) -> (x1,y1,w,h) for public_field.py
|
|
det[:, 2] = det[:, 2] - det[:, 0]
|
|
det[:, 3] = det[:, 3] - det[:, 1]
|
|
# classes(0~79) -> classes(1~80) for public_field.py
|
|
if e2e_coco:
|
|
classes_mapping(det, num_classes)
|
|
det = det.cpu().numpy()
|
|
dets.append(det)
|
|
|
|
if dets and len(dets) > 0:
|
|
dets = np.asarray(dets)
|
|
dets = np.squeeze(dets, axis=0) # remove outer []
|
|
dets = dets.tolist()
|
|
|
|
return dets
|
|
|
|
def make_grid(nx=20, ny=20):
|
|
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
|
grids = torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
|
|
return grids
|
|
|
|
def Yolov5_postprocess_onnx_sig(out,img_shape, im0_shape, conf_thres, iou_thres, top_k_num, grids, num_classes, anchors,vanish_point, e2e_coco) :
|
|
nc = num_classes # number of classes
|
|
no = nc + 5 # number of outputs per anchor
|
|
nl = len(anchors) # number of detection layers
|
|
na = len(anchors[0]) // 2 # number of anchors
|
|
a = torch.tensor(anchors).float().view(3, -1, 2)
|
|
anchor_grid = a.clone().view(3, 1, -1, 1, 1, 2)
|
|
stride = torch.tensor([ 8., 16., 32.])
|
|
z = []
|
|
for i in range(nl):
|
|
x = torch.from_numpy(out[i])
|
|
# print('x.shape',x.shape)
|
|
bs, _, ny, nx = x.shape # x(bs,3,20,20,85)
|
|
x = x.view(bs, na, no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
|
# grid_r = make_grid(nx, ny) ##grid
|
|
# grid_r = grid_r.numpy() ##grid
|
|
# file_name = str(i)+'.npy' ##grid
|
|
# np.save(file_name,grid_r) ##grid
|
|
grid = grids[i]#
|
|
#y = x.sigmoid()
|
|
y = x
|
|
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid) * stride[i] # xy
|
|
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i] # wh
|
|
z.append(y.view(bs, -1, no))
|
|
|
|
pred = torch.cat(z, 1)
|
|
return Yolov5_postprocess(pred, img_shape, im0_shape, conf_thres, iou_thres, top_k_num, num_classes,vanish_point, e2e_coco)
|
|
|
|
def Yolov5_postprocess_sig(out,img_shape, im0_shape, conf_thres, iou_thres, top_k_num, grids, num_classes, anchors,vanish_point, e2e_coco) :
|
|
nc = num_classes # number of classes
|
|
no = nc + 5 # number of outputs per anchor
|
|
nl = len(anchors) # number of detection layers
|
|
na = len(anchors[0]) // 2 # number of anchors
|
|
a = torch.tensor(anchors).float().view(3, -1, 2)
|
|
anchor_grid = a.clone().view(3, 1, -1, 1, 1, 2).to(out[0].device)
|
|
stride = torch.tensor([ 8., 16., 32.]).to(out[0].device)
|
|
z = []
|
|
for i in range(nl):
|
|
x = out[i]
|
|
bs, _, ny, nx = x.shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
|
# print('x.shape',x.shape)
|
|
x = x.view(bs, na, no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
|
|
|
# grid_r = make_grid(nx, ny) ##grid
|
|
# grid_r = grid_r.numpy() ##grid
|
|
# file_name = str(i)+'.npy' ##grid
|
|
# np.save(file_name,grid_r) ##grid
|
|
|
|
grid = grids[i].to(out[0].device) #
|
|
#y = x.sigmoid()
|
|
y = x
|
|
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid) * stride[i] # xy
|
|
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i] # wh
|
|
z.append(y.view(bs, -1, no))
|
|
# exit(0)
|
|
pred = torch.cat(z, 1)
|
|
return Yolov5_postprocess(pred, img_shape, im0_shape, conf_thres, iou_thres, top_k_num, num_classes,vanish_point, e2e_coco) |