|
import os |
|
import argparse |
|
import random |
|
import onnxruntime |
|
import numpy as np |
|
import cv2 |
|
from PIL import Image |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
|
|
from utils import input_transform, pad_image, resize_image, preprocess, get_confusion_matrix |
|
|
|
parser = argparse.ArgumentParser(description='HRNet') |
|
parser.add_argument('-m', '--onnx-model', default='', |
|
type=str, help='Path to onnx model.') |
|
parser.add_argument('-idir', '--img-dir', default='', |
|
type=str, help='Path to image filehold.') |
|
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 seting provider_options.") |
|
|
|
args = parser.parse_args() |
|
|
|
INPUT_SIZE = [512, 1024] |
|
|
|
|
|
def run_onnx_inference(ort_session, img): |
|
"""Infer an image with onnx seession |
|
|
|
Args: |
|
ort_session: Onnx session |
|
img (ndarray): Image to be infered. |
|
|
|
Returns: |
|
ndarray: Model inference result. |
|
""" |
|
pre_img, pad_h, pad_w = preprocess(img) |
|
img = np.expand_dims(pre_img, 0) |
|
ort_inputs = {ort_session.get_inputs()[0].name: img} |
|
o1 = ort_session.run(None, ort_inputs)[0] |
|
h, w = o1.shape[-2:] |
|
h_cut = int(h / INPUT_SIZE[0] * pad_h) |
|
w_cut = int(w / INPUT_SIZE[1] * pad_w) |
|
o1 = o1[..., :h - h_cut, :w - w_cut] |
|
return o1 |
|
|
|
|
|
def vis(out, image, save_path='color_.png'): |
|
|
|
pallete = [128, 64, 128, 244, 35, 232, 70, 70, 70, 102, 102, 156, 190, 153, 153, 153, 153, 153, 250, 170, 30, |
|
220, 220, 0, 107, 142, 35, 152, 251, 152, 70, 130, 180, 220, 20, 60, 255, 0, 0, 0, 0, 142, 0, 0, 70, |
|
0, 60, 100, 0, 80, 100, 0, 0, 230, 119, 11, 32 ] |
|
|
|
if out.shape[2] != image.shape[0] or out.shape[3] != image.shape[1]: |
|
out = torch.from_numpy(out).cpu() |
|
out = F.interpolate( |
|
out, size=image.shape[:2], |
|
mode='bilinear' |
|
).numpy() |
|
classMap_numpy = np.argmax(out[0], axis=0) |
|
classMap_numpy = Image.fromarray(classMap_numpy.astype(np.uint8)) |
|
classMap_numpy_color = classMap_numpy.copy() |
|
classMap_numpy_color.putpalette(pallete) |
|
classMap_numpy_color.save(save_path) |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
onnx_path = args.onnx_model |
|
if args.ipu: |
|
providers = ["VitisAIExecutionProvider"] |
|
provider_options = [{"config_file": args.provider_config}] |
|
else: |
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] |
|
provider_options = None |
|
|
|
img_dir = args.img_dir |
|
ort_session = onnxruntime.InferenceSession(onnx_path, providers=providers, provider_options=provider_options) |
|
|
|
img_names = os.listdir(img_dir) |
|
for img_name in img_names: |
|
image_path = os.path.join(img_dir, img_name) |
|
img = cv2.imread(image_path) |
|
img_vis = np.copy(img) |
|
outs = run_onnx_inference(ort_session, img) |
|
vis(outs, img_vis) |
|
|
