File size: 3,061 Bytes
c2ba1ac c139485 c2ba1ac c139485 c2ba1ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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)
img = np.transpose(img, (0,2,3,1))
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 ]
# out = out[0]
out = out.transpose(0, 3, 1, 2)
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)
|