TDA4VM: Inference with Custom Artifacts Kills Kernel in Edge AI Studio

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import os
import re
import sys
import cv2
import tqdm
import onnx
import math
import copy
import shutil
import platform
import itertools

import numpy as np
import onnxruntime as rt
import ipywidgets as widgets
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from pathlib import Path
from munkres import Munkres
from numpy.lib.stride_tricks import as_strided
from IPython.display import Markdown as md
from PIL import Image, ImageFont, ImageDraw, ImageEnhance
from scripts.utils import imagenet_class_to_name, download_model, loggerWritter, get_svg_path, get_preproc_props, single_img_visualise


# In[2]:


def preprocess_for_onnx_pose_estimation(image_path, size, mean, scale, layout, reverse_channels, pad_color=114, pad_type="center"):
    # Step 1
    # read the image using openCVimport json_tricks as json
    img = cv2.imread(image_path)
    
    # Step 2
    # convert to RGB
    img = img[:,:,::-1]
    
    # Step 3    
    # Most of the onnx models are trained using
    # 512x512 images. The general rule of thumb
    # is to scale the input image while preserving
    # the original aspect ratio so that the
    # longer edge is 512 pixels, and then
    # pad the scaled image to 512x512
    
    size = (size,size) if not isinstance(size, (list,tuple)) else size
    desired_size = size[-1]
    old_size = img.shape[:2] # old_size is in (height, width) format

    ratio = float(desired_size)/max(old_size)
    new_size = tuple([int(x*ratio) for x in old_size])

    # new_size should be in (width, height) format
    img = cv2.resize(img, (new_size[1], new_size[0]))

    delta_w = size[1] - new_size[1]
    delta_h = size[0] - new_size[0]

    if pad_type=="corner":
        top, left = 0, 0
        bottom, right = delta_h, delta_w
    else:
        delta_w = size[1] - new_size[1]
        delta_h = size[0] - new_size[0]
        top, bottom = delta_h//2, delta_h-(delta_h//2)
        left, right = delta_w//2, delta_w-(delta_w//2)


    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT,
        value=pad_color)
    
    # Step 4
    # Apply scaling and mean subtraction.
    # if your model is built with an input
    # normalization layer, then you might
    # need to skip this
    if mean is not None and scale is not None:
        img = img.astype('float32')
        for mean, scale, ch in zip(mean, scale, range(img.shape[2])):
            img[:,:,ch] = ((img.astype('float32')[:,:,ch] - mean) * scale)
            
    # Step 5
    if reverse_channels:
        img = img[:,:,::-1]
    
    # Step 6
    img = np.expand_dims(img,axis=0)
    img = np.transpose(img, (0, 3, 1, 2))
    
    return img, top, left, ratio


# In[3]:


calib_images = [
'sample-images/yoga0.jpg',
'sample-images/yoga1.bmp',
'sample-images/yoga2.bmp',
]

output_dir = '../custom-artifacts-temp/onnx/yolox_s_pose_ti_lite_49p5_78p0.onnx'
#onnx_model_path_TDA4VM = '/opt/model_zoo/ONR-KD-7060-human-pose-yolox-s-640x640/model/yolox_s_pose_ti_lite_49p5_78p0.onnx'
#onnx_model_path_EdgeAIcloud = '/home/root/notebooks/model-zoo/models/vision/keypoint/coco/edgeai-yolox/yolox_s_pose_ti_lite_49p5_78p0.onnx'
onnx_model_path_EdgeAIcloud = '/home/root/notebooks/prebuilt-models/8bits/kd-7060_onnxrt_coco_edgeai-yolox_yolox_s_pose_ti_lite_49p5_78p0_onnx/model/yolox_s_pose_ti_lite_49p5_78p0.onnx'
onnx.shape_inference.infer_shapes_path(onnx_model_path_EdgeAIcloud, onnx_model_path_EdgeAIcloud)

#compilation options - knobs to tweak 
num_bits =8
accuracy =1

log_dir = Path("logs").mkdir(parents=True, exist_ok=True)

# stdout and stderr saved to a *.log file.  
with loggerWritter("logs/custon-model-onnx"):
    
    # model compilation options
    compile_options = {
        'tidl_tools_path' : os.environ['TIDL_TOOLS_PATH'],
        'artifacts_folder' : output_dir,
        'tensor_bits' : num_bits,
        'accuracy_level' : accuracy,
        'advanced_options:calibration_frames' : len(calib_images), 
        'advanced_options:calibration_iterations' : 3, # used if accuracy_level = 1
        'debug_level' : 1,
        'deny_list' : "MaxPool" #Comma separated string of operator types as defined by ONNX runtime, ex "MaxPool, Concat"
    }

# create the output dir if not present
# clear the directory
os.makedirs(output_dir, exist_ok=True)
for root, dirs, files in os.walk(output_dir, topdown=False):
    [os.remove(os.path.join(root, f)) for f in files]
    [os.rmdir(os.path.join(root, d)) for d in dirs]


# In[4]:


# create & compile model with compile options specified above 
so = rt.SessionOptions()
EP_list = ['TIDLCompilationProvider','CPUExecutionProvider']
sess = rt.InferenceSession(onnx_model_path_EdgeAIcloud ,providers=EP_list, provider_options=[compile_options, {}], sess_options=so)

input_details = sess.get_inputs()
print('input_details: ', input_details)

output_details = sess.get_outputs()
print('output_details: ', output_details)


label = 'ONR-KD-7060-human-pose-yolox-s-640x640'
pad_color = 128 if 'ae' in label and 'yolo' not in label else 114
pad_type = "corner" if 'yolox' in label else "center"
size = 640
mean = [0.0, 0.0, 0.0]
scale = [1.0, 1.0, 1.0]
layout = 0
reverse_channels = True

# run inference for each calibration image 
for num in tqdm.trange(len(calib_images)):
    #output = list(sess.run(None, {input_details[0].name : preprocess_for_onnx_pose_estimation(calib_images[num], size, mean, scale, layout, reverse_channels, pad_color, pad_type)}))[0]
    image_name = calib_images[num]
    print('label = ', label)
    print('pad_color = ', pad_color)
    print('pad_type = ', pad_type)
    print('image_name = ', image_name)
    processed_image, top, left, ratio = preprocess_for_onnx_pose_estimation(image_name, size, mean, scale, layout, reverse_channels, pad_color, pad_type)
    
    print('processed_image', processed_image)
    print('top', top)  
    print('left', left)
    print('ratio', ratio)
    
    if not input_details[0].type == 'tensor(float)':
        processed_image = np.uint8(processed_image)

    image_size = processed_image.shape[3]    
    print('image_size = ', image_size)
    out_file=None
    output=None
    output = list(sess.run(None, {input_details[0].name : processed_image})) #[0]
    print('output = ', output)
    
    #%matplotlib inline
    #output_image = single_img_visualise(output, image_size, calib_images[num], out_file, top, left, ratio, udp=True, thickness=2, radius=5, label=label)

    # plot the outut using matplotlib
    #plt.rcParams["figure.figsize"]=20,20
    #plt.rcParams['figure.dpi'] = 200 # 200 e.g. is really fine, but slower
    #plt.imshow(output_image)
    #plt.show()


# In[5]:


# subgraphs visualization for debugging
subgraph_link =get_svg_path(output_dir) 
for sg in subgraph_link:
    hl_text = os.path.join(*Path(sg).parts[4:])
    sg_rel = os.path.join('../', sg)
    display(md("[{}]({})".format(hl_text,sg_rel)))


# In[6]:


root_src_dir = output_dir
root_dst_dir = 'custom-artifacts/onnx/yolox_s_pose_ti_lite_49p5_78p0.onnx'

for src_dir, dirs, files in os.walk(root_src_dir):
    dst_dir = src_dir.replace(root_src_dir, root_dst_dir, 1)
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    for file_ in files:
        src_file = os.path.join(src_dir, file_)
        dst_file = os.path.join(dst_dir, file_)
        if os.path.exists(dst_file):
            os.remove(dst_file)
        shutil.copy(src_file, dst_dir)


# In[7]:


del so
print('closed calibration inference session...')


# In[8]:


# use compiled image for inference 
out_file=None
image_name = 'sample-images/yoga0.jpg'
delegate_options = {
    'artifacts_folder': './custom-artifacts/onnx/yolox_s_pose_ti_lite_49p5_78p0.onnx'
}
print('delegate_options: ', delegate_options)


# In[9]:


so0 = rt.SessionOptions()
EP_list = ['TIDLCompilationProvider','CPUExecutionProvider']
print('EP_list: ', EP_list)


# In[ ]:


sess0 = rt.InferenceSession(onnx_model_path_EdgeAIcloud ,providers=EP_list, provider_options=[delegate_options, {}], sess_options=so0)
print('session0 started')


# In[ ]:


input_details0 = sess0.get_inputs()
print('input_details0: ', input_details0)


# In[ ]:


processed_image, top, left, ratio = preprocess_for_onnx_pose_estimation(image_name, size, mean, scale, layout, reverse_channels, pad_color, pad_type)

if not input_details[0].type == 'tensor(float)':
    processed_image = np.uint8(processed_image)

image_size = processed_image.shape[3]    
output0 = list(sess0.run(None, {input_details0[0].name : processed_image}))[0]


# In[ ]:


# post processing 
get_ipython().run_line_magic('matplotlib', 'inline')
output_image = single_img_visualise(output0, image_size, image_name, out_file, top, left, ratio, udp=True, thickness=2, radius=5, label=label)

# plot the outut using matplotlib
plt.rcParams["figure.figsize"]=20,20
plt.rcParams['figure.dpi'] = 200 # 200 e.g. is really fine, but slower
plt.imshow(output_image)
plt.show()