|
import torch |
|
import torch.nn as nn |
|
import onnxruntime |
|
import numpy as np |
|
import argparse |
|
from utils import ( |
|
LoadImages, |
|
non_max_suppression, |
|
plot_images, |
|
output_to_target, |
|
) |
|
import sys |
|
import pathlib |
|
CURRENT_DIR = pathlib.Path(__file__).parent |
|
sys.path.append(str(CURRENT_DIR)) |
|
|
|
def preprocess(img): |
|
img = torch.from_numpy(img) |
|
img = img.float() |
|
img /= 255 |
|
return img |
|
|
|
|
|
class DFL(nn.Module): |
|
|
|
def __init__(self, c1=16): |
|
super().__init__() |
|
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False) |
|
x = torch.arange(c1, dtype=torch.float) |
|
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1)) |
|
self.c1 = c1 |
|
|
|
def forward(self, x): |
|
b, c, a = x.shape |
|
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view( |
|
b, 4, a |
|
) |
|
|
|
|
|
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) |
|
return torch.cat((x1y1, x2y2), dim) |
|
|
|
|
|
def post_process(x): |
|
dfl = DFL(16) |
|
anchors = torch.tensor( |
|
np.load( |
|
"./anchors.npy", |
|
allow_pickle=True, |
|
) |
|
) |
|
strides = torch.tensor( |
|
np.load( |
|
"./strides.npy", |
|
allow_pickle=True, |
|
) |
|
) |
|
box, cls = torch.cat([xi.view(x[0].shape[0], 144, -1) for xi in x], 2).split( |
|
(16 * 4, 80), 1 |
|
) |
|
dbox = dist2bbox(dfl(box), anchors.unsqueeze(0), xywh=True, dim=1) * strides |
|
y = torch.cat((dbox, cls.sigmoid()), 1) |
|
return y, x |
|
|
|
|
|
def make_parser(): |
|
parser = argparse.ArgumentParser("onnxruntime inference sample") |
|
parser.add_argument( |
|
"-m", |
|
"--onnx_model", |
|
type=str, |
|
default="./yolov8m.onnx", |
|
help="input your onnx model.", |
|
) |
|
parser.add_argument( |
|
"-i", |
|
"--image_path", |
|
type=str, |
|
default='./demo.jpg', |
|
help="path to your input image.", |
|
) |
|
parser.add_argument( |
|
"-o", |
|
"--output_path", |
|
type=str, |
|
default='./demo_infer.jpg', |
|
help="path to your output directory.", |
|
) |
|
parser.add_argument( |
|
"--ipu", action='store_true', help='flag for ryzen ai' |
|
) |
|
parser.add_argument( |
|
"--provider_config", default='', type=str, help='provider config for ryzen ai' |
|
) |
|
return parser |
|
|
|
classnames = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', |
|
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', |
|
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', |
|
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', |
|
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', |
|
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', |
|
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', |
|
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', |
|
'hair drier', 'toothbrush'] |
|
names = {k: classnames[k] for k in range(80)} |
|
imgsz = [640, 640] |
|
|
|
|
|
if __name__ == '__main__': |
|
args = make_parser().parse_args() |
|
source = args.image_path |
|
dataset = LoadImages( |
|
source, imgsz=imgsz, stride=32, auto=False, transforms=None, vid_stride=1 |
|
) |
|
onnx_weight = args.onnx_model |
|
if args.ipu: |
|
providers = ["VitisAIExecutionProvider"] |
|
provider_options = [{"config_file": args.provider_config}] |
|
onnx_model = onnxruntime.InferenceSession(onnx_weight, providers=providers, provider_options=provider_options) |
|
else: |
|
onnx_model = onnxruntime.InferenceSession(onnx_weight) |
|
for batch in dataset: |
|
path, im, im0s, vid_cap, s = batch |
|
im = preprocess(im) |
|
if len(im.shape) == 3: |
|
im = im[None] |
|
|
|
|
|
outputs = onnx_model.run(None, {onnx_model.get_inputs()[0].name: im.permute(0, 2, 3, 1).cpu().numpy()}) |
|
outputs = [torch.tensor(item).permute(0, 3, 1, 2) for item in outputs] |
|
preds = post_process(outputs) |
|
preds = non_max_suppression( |
|
preds, 0.25, 0.7, agnostic=False, max_det=300, classes=None |
|
) |
|
plot_images( |
|
im, |
|
*output_to_target(preds, max_det=15), |
|
source, |
|
fname=args.output_path, |
|
names=names, |
|
) |
|
|
