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