| | 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.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) |
| |
|
