| import os | |
| import sys | |
| import pathlib | |
| CURRENT_DIR = pathlib.Path(__file__).parent | |
| sys.path.append(str(CURRENT_DIR)) | |
| import numpy as np | |
| from tqdm import tqdm | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils import data | |
| import torchvision.transforms as transform | |
| import torch.nn.functional as F | |
| import onnxruntime | |
| from PIL import Image | |
| import argparse | |
| from datasets.utils import colorize_mask, build_img | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description='SemanticFPN model') | |
| parser.add_argument('--onnx_path', type=str, default='FPN_int_NHWC.onnx') | |
| parser.add_argument('--save_path', type=str, default='./data/demo_results/senmatic_results.png') | |
| parser.add_argument('--input_path', type=str, default='data/cityscapes/leftImg8bit/test/bonn/bonn_000000_000019_leftImg8bit.png') | |
| parser.add_argument('--ipu', action='store_true', help='use ipu') | |
| parser.add_argument('--provider_config', type=str, default=None, | |
| help='provider config path') | |
| args = parser.parse_args() | |
| if args.ipu: | |
| providers = ["VitisAIExecutionProvider"] | |
| provider_options = [{"config_file": args.provider_config}] | |
| else: | |
| providers = ['CPUExecutionProvider'] | |
| provider_options = None | |
| onnx_path = args.onnx_path | |
| input_img = build_img(args) | |
| session = onnxruntime.InferenceSession(onnx_path, providers=providers, provider_options=provider_options) | |
| ort_input = {session.get_inputs()[0].name: input_img.cpu().numpy().transpose(0,2,3,1)} | |
| ort_output = session.run(None, ort_input)[0].transpose(0,3,1,2) | |
| if isinstance(ort_output, (tuple, list)): | |
| ort_output = ort_output[0] | |
| output = ort_output[0].transpose(1, 2, 0) | |
| seg_pred = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) | |
| color_mask = colorize_mask(seg_pred) | |
| os.makedirs(os.path.dirname(args.save_path), exist_ok=True) | |
| color_mask.save(args.save_path) | |