|
|
|
|
|
|
|
import os |
|
import sys |
|
import cv2 |
|
import pathlib |
|
import argparse |
|
import numpy as np |
|
import onnxruntime as ort |
|
|
|
CURRENT_DIR = pathlib.Path(__file__).parent |
|
sys.path.append(str(CURRENT_DIR)) |
|
|
|
from coco import COCO_CLASSES |
|
from demo_utils import mkdir, multiclass_nms, demo_postprocess, vis |
|
|
|
|
|
def make_parser(): |
|
parser = argparse.ArgumentParser("onnxruntime inference sample") |
|
parser.add_argument( |
|
"-m", |
|
"--model", |
|
type=str, |
|
default="yolox-s-int8.onnx", |
|
help="Input your onnx model.", |
|
) |
|
parser.add_argument( |
|
"-i", |
|
"--image_path", |
|
type=str, |
|
default='test_image.png', |
|
help="Path to your input image.", |
|
) |
|
parser.add_argument( |
|
"-o", |
|
"--output_dir", |
|
type=str, |
|
default='demo_output', |
|
help="Path to your output directory.", |
|
) |
|
parser.add_argument( |
|
"-s", |
|
"--score_thr", |
|
type=float, |
|
default=0.3, |
|
help="Score threshold to filter the result.", |
|
) |
|
parser.add_argument( |
|
"--input_shape", |
|
type=str, |
|
default="640,640", |
|
help="Specify an input shape for inference.", |
|
) |
|
parser.add_argument( |
|
"--ipu", |
|
action="store_true", |
|
help="Use IPU for inference.", |
|
) |
|
parser.add_argument( |
|
"--provider_config", |
|
type=str, |
|
default="vaip_config.json", |
|
help="Path of the config file for setting provider_options.", |
|
) |
|
return parser |
|
|
|
|
|
def preprocess(img, input_shape, swap=(2, 0, 1)): |
|
""" |
|
Preprocessing part of YOLOX for scaling and padding image as input to the network. |
|
|
|
Args: |
|
img (numpy.ndarray): H x W x C, image read with OpenCV |
|
input_shape (tuple(int)): input shape of the network for inference |
|
swap (tuple(int)): new order of axes to transpose the input image |
|
|
|
Returns: |
|
padded_img (numpy.ndarray): preprocessed image to be fed to the network |
|
ratio (float): ratio for scaling the image to the input shape |
|
""" |
|
if len(img.shape) == 3: |
|
padded_img = np.ones((input_shape[0], input_shape[1], 3), dtype=np.uint8) * 114 |
|
else: |
|
padded_img = np.ones(input_shape, dtype=np.uint8) * 114 |
|
ratio = min(input_shape[0] / img.shape[0], input_shape[1] / img.shape[1]) |
|
resized_img = cv2.resize( |
|
img, |
|
(int(img.shape[1] * ratio), int(img.shape[0] * ratio)), |
|
interpolation=cv2.INTER_LINEAR, |
|
).astype(np.uint8) |
|
padded_img[: int(img.shape[0] * ratio), : int(img.shape[1] * ratio)] = resized_img |
|
padded_img = padded_img.transpose(swap) |
|
padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) |
|
return padded_img, ratio |
|
|
|
|
|
def postprocess(outputs, input_shape, ratio): |
|
""" |
|
Post-processing part of YOLOX for generating final results from outputs of the network. |
|
|
|
Args: |
|
outputs (tuple(numpy.ndarray)): outputs of the detection heads with onnxruntime session |
|
input_shape (tuple(int)): input shape of the network for inference |
|
ratio (float): ratio for scaling the image to the input shape |
|
|
|
Returns: |
|
dets (numpy.ndarray): n x 6, dets[:,:4] -> boxes, dets[:,4] -> scores, dets[:,5] -> class indices |
|
""" |
|
outputs = [out.reshape(*out.shape[:2], -1).transpose(0,2,1) for out in outputs] |
|
outputs = np.concatenate(outputs, axis=1) |
|
outputs[..., 4:] = sigmoid(outputs[..., 4:]) |
|
predictions = demo_postprocess(outputs, input_shape, p6=False)[0] |
|
boxes = predictions[:, :4] |
|
scores = predictions[:, 4:5] * predictions[:, 5:] |
|
boxes_xyxy = np.ones_like(boxes) |
|
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. |
|
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. |
|
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. |
|
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. |
|
boxes_xyxy /= ratio |
|
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) |
|
return dets |
|
|
|
|
|
def sigmoid(x): |
|
return 1.0 / (1.0 + np.exp(-x)) |
|
|
|
|
|
if __name__ == '__main__': |
|
args = make_parser().parse_args() |
|
input_shape = tuple(map(int, args.input_shape.split(','))) |
|
origin_img = cv2.imread(args.image_path) |
|
img, ratio = preprocess(origin_img, input_shape) |
|
if args.ipu: |
|
providers = ["VitisAIExecutionProvider"] |
|
provider_options = [{"config_file": args.provider_config}] |
|
else: |
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] |
|
provider_options = None |
|
session = ort.InferenceSession(args.model, providers=providers, provider_options=provider_options) |
|
|
|
ort_inputs = {session.get_inputs()[0].name: np.transpose(img[None, :, :, :], (0, 2 ,3, 1))} |
|
outputs = session.run(None, ort_inputs) |
|
outputs = [np.transpose(out, (0, 3, 1, 2)) for out in outputs] |
|
dets = postprocess(outputs, input_shape, ratio) |
|
if dets is not None: |
|
final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] |
|
origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds, |
|
conf=args.score_thr, class_names=COCO_CLASSES) |
|
mkdir(args.output_dir) |
|
output_path = os.path.join(args.output_dir, os.path.basename(args.image_path)) |
|
cv2.imwrite(output_path, origin_img) |
|
|
