| | import os |
| | import argparse |
| | import random |
| | import onnxruntime |
| | import numpy as np |
| |
|
| | import torch |
| | from torch.nn import functional as F |
| | from torch.utils import data |
| |
|
| | import cv2 |
| | from PIL import Image |
| | from tqdm import tqdm |
| |
|
| | 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('-r', '--root', default='', |
| | type=str, help='Path to dataset root.') |
| | parser.add_argument('-l', '--list_path', default='', |
| | type=str, help='Path to dataset list.') |
| | 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] |
| | NUM_CLASSES = 19 |
| | IGNORE_LABEL = 255 |
| |
|
| |
|
| | class Cityscapes(data.Dataset): |
| | def __init__(self, |
| | root, |
| | list_path, |
| | num_classes=19, |
| | downsample_rate=8, |
| | ignore_label=-1): |
| |
|
| | self.root = root |
| | self.list_path = list_path |
| | self.num_classes = num_classes |
| | self.downsample_rate = downsample_rate |
| | |
| | self.img_list = [line.strip().split() for line in open(root+list_path)] |
| |
|
| | self.files = self.read_files() |
| |
|
| | self.label_mapping = {-1: ignore_label, 0: ignore_label, |
| | 1: ignore_label, 2: ignore_label, |
| | 3: ignore_label, 4: ignore_label, |
| | 5: ignore_label, 6: ignore_label, |
| | 7: 0, 8: 1, 9: ignore_label, |
| | 10: ignore_label, 11: 2, 12: 3, |
| | 13: 4, 14: ignore_label, 15: ignore_label, |
| | 16: ignore_label, 17: 5, 18: ignore_label, |
| | 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, |
| | 25: 12, 26: 13, 27: 14, 28: 15, |
| | 29: ignore_label, 30: ignore_label, |
| | 31: 16, 32: 17, 33: 18} |
| | |
| | def read_files(self): |
| | files = [] |
| | for item in self.img_list: |
| | image_path, label_path = item |
| | name = os.path.splitext(os.path.basename(label_path))[0] |
| | files.append({ |
| | "img": image_path, |
| | "label": label_path, |
| | "name": name, |
| | }) |
| | return files |
| |
|
| | def __len__(self): |
| | return len(self.files) |
| | |
| | def convert_label(self, label, inverse=False): |
| | temp = label.copy() |
| | if inverse: |
| | for v, k in self.label_mapping.items(): |
| | label[temp == k] = v |
| | else: |
| | for k, v in self.label_mapping.items(): |
| | label[temp == k] = v |
| | return label |
| |
|
| | def __getitem__(self, index): |
| | item = self.files[index] |
| | image = cv2.imread(os.path.join(self.root, item["img"]), |
| | cv2.IMREAD_COLOR) |
| | label = cv2.imread(os.path.join(self.root, item["label"]), |
| | cv2.IMREAD_GRAYSCALE) |
| | label = self.convert_label(label) |
| | image, label = self.gen_sample(image, label) |
| |
|
| | return image.copy(), label.copy() |
| |
|
| | def gen_sample(self, image, label): |
| | label = self.label_transform(label) |
| | |
| |
|
| | if self.downsample_rate != 1: |
| | label = cv2.resize( |
| | label, |
| | None, |
| | fx=self.downsample_rate, |
| | fy=self.downsample_rate, |
| | interpolation=cv2.INTER_NEAREST |
| | ) |
| |
|
| | return image, label |
| |
|
| | def label_transform(self, label): |
| | return np.array(label).astype('int32') |
| |
|
| |
|
| | 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 testval(ort_session, root, list_path): |
| |
|
| | test_dataset = Cityscapes( |
| | root=root, |
| | list_path=list_path, |
| | num_classes=NUM_CLASSES, |
| | ignore_label=IGNORE_LABEL, |
| | downsample_rate=1) |
| |
|
| | testloader = torch.utils.data.DataLoader( |
| | test_dataset, |
| | batch_size=1, |
| | shuffle=False, |
| | num_workers=4, |
| | pin_memory=True) |
| |
|
| | confusion_matrix = np.zeros( |
| | (NUM_CLASSES, NUM_CLASSES)) |
| | for index, batch in enumerate(tqdm(testloader)): |
| | image, label = batch |
| | image = image.numpy()[0] |
| | out = run_onnx_inference(ort_session, image) |
| | size = label.size() |
| | |
| | out = out.transpose(0, 3, 1, 2) |
| | if out.shape[2] != size[1] or out.shape[3] != size[2]: |
| | out = torch.from_numpy(out).cpu() |
| | pred = F.interpolate( |
| | out, size=size[1:], |
| | mode='bilinear' |
| | ) |
| |
|
| | confusion_matrix += get_confusion_matrix( |
| | label, |
| | pred, |
| | size, |
| | NUM_CLASSES, |
| | IGNORE_LABEL) |
| |
|
| | pos = confusion_matrix.sum(1) |
| | res = confusion_matrix.sum(0) |
| | tp = np.diag(confusion_matrix) |
| | pixel_acc = tp.sum()/pos.sum() |
| | mean_acc = (tp/np.maximum(1.0, pos)).mean() |
| | IoU_array = (tp / np.maximum(1.0, pos + res - tp)) |
| | mean_IoU = IoU_array.mean() |
| |
|
| | return mean_IoU, IoU_array, pixel_acc, mean_acc |
| |
|
| |
|
| | if __name__ == "__main__": |
| |
|
| | onnx_path = args.onnx_model |
| | root = args.root |
| | list_path = args.list_path |
| | if args.ipu: |
| | providers = ["VitisAIExecutionProvider"] |
| | provider_options = [{"config_file": args.provider_config}] |
| | else: |
| | providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] |
| | provider_options = None |
| | |
| | ort_session = onnxruntime.InferenceSession(onnx_path, providers=providers, provider_options=provider_options) |
| | |
| | mean_IoU, IoU_array, pixel_acc, mean_acc = testval(ort_session, root, list_path) |
| |
|
| | msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, Mean_Acc: {: 4.4f}'.format(mean_IoU, \ |
| | pixel_acc, mean_acc) |
| | print(msg) |
| |
|
