Skip to content

Commit

Permalink
moved v8 nms to dedicated function for better readibility, adapted v8…
Browse files Browse the repository at this point in the history 
… nms logic to match code used in https://github.com/ultralytics/ultralytics.git/examples/YOLOv8-ONNXRuntime/main.py
  • Loading branch information
@sebastianfis
sebastianfis committed Mar 24, 2024
1 parent 18f6fa1 commit 821deb9
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 25 deletions.
18 changes: 11 additions & 7 deletions edgetpumodel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import numpy as np
import pycoral.utils.edgetpu as etpu
from pycoral.adapters import common
from nms import non_max_suppression
from nms import non_max_suppression, non_max_suppresion_v8
import cv2
import json

Expand Down Expand Up @@ -42,7 +42,7 @@ def __init__(self, model_file, names_file, conf_thresh=0.25, iou_thresh=0.45, fi
self.iou_thresh = iou_thresh
self.filter_classes = filter_classes
self.agnostic_nms = agnostic_nms
self.max_det = 1000
self.max_det = max_det
self.v8 = v8

logger.info("Confidence threshold: {}".format(conf_thresh))
Expand Down Expand Up @@ -150,7 +150,7 @@ def forward(self, x:np.ndarray, with_nms=True) -> np.ndarray:
tstart = time.time()
# Transpose if C, H, W
if x.shape[0] == 3:
x = x.transpose((1,2,0))
x = x.transpose((1, 2, 0))

x = x.astype('float32')

Expand All @@ -167,15 +167,19 @@ def forward(self, x:np.ndarray, with_nms=True) -> np.ndarray:
# Scale output
result = (common.output_tensor(self.interpreter, 0).astype('float32') - self.output_zero) * self.output_scale
if self.v8:
result = np.transpose(result, [0, 2, 1]) # tranpose for yoolov8 models
result = np.transpose(result, [0, 2, 1]) # tranpose for yolov8 models

self.inference_time = time.time() - tstart

if with_nms:

tstart = time.time()
nms_result = non_max_suppression(result, self.conf_thresh, self.iou_thresh, self.filter_classes,
self.agnostic_nms, max_det=self.max_det, v8=self.v8)
if self.v8:
nms_result = non_max_suppresion_v8(result, self.conf_thresh, self.iou_thresh, self.filter_classes,
self.agnostic_nms, max_det=self.max_det)
else:
nms_result = non_max_suppression(result, self.conf_thresh, self.iou_thresh, self.filter_classes,
self.agnostic_nms, max_det=self.max_det)
self.nms_time = time.time() - tstart

return nms_result
Expand Down Expand Up @@ -238,7 +242,7 @@ def process_predictions(self, det, output_image, pad, output_path="detection.jpg
if len(det):
# Rescale boxes from img_size to im0 size
# x1, y1, x2, y2=
det[:, :4] = self.get_scaled_coords(det[:,:4], output_image, pad)
det[:, :4] = self.get_scaled_coords(det[:, :4], output_image, pad)
output = {}
base, ext = os.path.splitext(output_path)

Expand Down
122 changes: 104 additions & 18 deletions nms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,16 +48,103 @@ def nms(dets, scores, thresh):

return np.array(keep)

def box_iou(box_i,boxes):
raise NotImplementedError('merged iou calculation not implemented')

def non_max_suppresion_v8(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
labels=(), max_det=300):
# TODO: Test this for changed parameter number in yolov8!!! All that could be detected with last commit were persons!

# 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, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
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 = [np.zeros((0, 6))] * 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

# Cat apriori labels if autolabelling
if labels and len(labels[xi]):
l = labels[xi]
v = np.zeros((len(l), nc + 4))
v[:, :4] = l[:, 1:4] # box
v[:, 4] = 1.0 # conf
v[range(len(l)), l[:, 0].long() + 4] = 1.0 # cls
x = np.concatenate((x, v), 0)

# If none remain process next image
if not x.shape[0]:
continue

# Compute conf already done in v8

# 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[:, 4:] > conf_thres).nonzero(as_tuple=False).T
x = np.concatenate((box[i], x[i, j + 4, None], j[:, None].astype(float)), axis=1)
else: # best class only
conf = np.amax(x[:, 4:], axis=1, keepdims=True)
j = np.argmax(x[:, 4:], axis=1).reshape(conf.shape)
x = np.concatenate((box, conf, j.astype(float)), axis=1)[conf.flatten() > conf_thres]

# Filter by class
if classes is not None:
raise NotImplementedError('objectivity parameter not implemented for yolov8 models!')

# 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

# 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
boxes, scores = x[:, :4], x[:, 4] # boxes, 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 < 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] = np.dot(weights, x[:, :4]).astype(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 non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
labels=(), max_det=300, v8=False):
# TODO: Correct this for changed parameter number in yolov8!!!
if v8:
split_val = -1
else:
split_val = 0

nc = prediction.shape[2] - 5 - split_val # number of classes
labels=(), max_det=300):

nc = prediction.shape[2] - 5 # number of classes
xc = prediction[..., 4] > conf_thres # candidates

# Checks
Expand All @@ -82,30 +169,29 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
# Cat apriori labels if autolabelling
if labels and len(labels[xi]):
l = labels[xi]
v = np.zeros((len(l), nc + 5 + split_val))
v[:, :4] = l[:, 1:(5 + split_val)] # box
v[:, 4 + split_val] = 1.0 # conf
v[range(len(l)), l[:, 0].long() + 5 + split_val] = 1.0 # cls
v = np.zeros((len(l), nc + 5))
v[:, :4] = l[:, 1:5] # box
v[:, 4] = 1.0 # conf
v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls
x = np.concatenate((x, v), 0)

# If none remain process next image
if not x.shape[0]:
continue

# Compute conf
if not v8:
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_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 + split_val):] > conf_thres).nonzero(as_tuple=False).T
x = np.concatenate((box[i], x[i, j + 5 + split_val, None], j[:, None].astype(float)), axis=1)
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
x = np.concatenate((box[i], x[i, j + 5, None], j[:, None].astype(float)), axis=1)
else: # best class only
conf = np.amax(x[:, 5 + split_val:], axis=1, keepdims=True)
j = np.argmax(x[:, 5 + split_val:], axis=1).reshape(conf.shape)
conf = np.amax(x[:, 5:], axis=1, keepdims=True)
j = np.argmax(x[:, 5:], axis=1).reshape(conf.shape)
x = np.concatenate((box, conf, j.astype(float)), axis=1)[conf.flatten() > conf_thres]

# Filter by class
Expand Down

0 comments on commit 821deb9

Please sign in to comment.