[FAQ] How do I get a valid param.yaml file for my TIDL deep learning model? The demos in edgeai-gst-apps throw errors on my model

#!/usr/bin/python3

import os
import yaml
import argparse
import copy
from collections.abc import Iterable  

DEFAULT_NUM_BOXES = 200

def update_preprocess(options_dict):
    
    preprocess = options_dict.get('preprocess')
    if preprocess:
        mean = preprocess.get('mean')
        scale = preprocess.get('scale')
        data_layout = preprocess.get('data_layout', 'NCHW')
        reverse_channels = preprocess.get('reverse_channels', False)
    else: 
        mean = None
        scale = None
        data_layout = 'NCHW' #default; tflite usually using NHWC, ONNX uses NCHW
        reverse_channels = False

    input_optimization = False
    if mean is None and scale is None:
        input_optimization = True
    elif isinstance(mean, Iterable) and isinstance(scale, Iterable) and \
        all([m == 0 for m in mean]) and all([s == 1 for  s in scale]):
        input_optimization = True
    elif mean == 0 and scale == 1: 
        input_optimization = True

    options_dict['session']['input_mean'] = copy.copy(mean)
    options_dict['session']['input_scale'] = copy.copy(scale)
    options_dict['session']['input_data_layout'] = copy.copy(data_layout)
    options_dict['session']['input_optimization'] = input_optimization
    options_dict['session']['reverse_channels'] = reverse_channels
    options_dict['preprocess']['reverse_channels'] = reverse_channels



    return options_dict

def ort_detail_to_dict(detail):
    new_detail = {
        'name': detail.name,
        'shape': detail.shape,
        'type': detail.type
    }
    return new_detail

def tfl_detail_to_dict(detail):
    shape = detail['shape'].tolist()
    new_detail = {
        'name': detail['name'],
        'shape': shape,
        'type': str(detail['dtype'])
    }
    return new_detail

def update_input_details(options_dict, model, ext):
    in_details_yaml = options_dict['session'].get('input_details', None)
    if in_details_yaml is None:
        options_dict['session']['input_details'] = []
        in_details_yaml = options_dict['session'].get('input_details', None)

    if ext == '.onnx':
        input_details = model.get_inputs()
        if len(in_details_yaml) == len(input_details): return options_dict

        for in_detail in input_details:
            new_detail = ort_detail_to_dict(in_detail)
            for i,s in enumerate(new_detail['shape']):
                if s is None:
                    #Note that failures here may come from any other unspecified dimensions
                    new_detail['shape'][i] = DEFAULT_NUM_BOXES

            in_details_yaml.append(new_detail)

    elif ext == '.tflite':
        input_details = model.get_input_details()
        if len(in_details_yaml) == len(input_details): return options_dict

        for in_detail in input_details:
            new_detail = tfl_detail_to_dict(in_detail)
            in_details_yaml.append(new_detail)

    return options_dict

def update_output_details(options_dict, model, ext):
    out_details_yaml = options_dict['session'].get('output_details', None)
    if out_details_yaml is None:
        options_dict['session']['output_details'] = []
        out_details_yaml = options_dict['session'].get('output_details', None)

    if ext == '.onnx':
        output_details = model.get_outputs()
        if len(out_details_yaml) == len(output_details): return options_dict

        for out_detail in output_details:
            new_detail = ort_detail_to_dict(out_detail)
            for i,s in enumerate(new_detail['shape']):
                if s is None:
                    #Note that failures here may come from any other unspecified dimensions
                    new_detail['shape'][i] = DEFAULT_NUM_BOXES

            out_details_yaml.append(new_detail)

    elif ext == '.tflite':
        output_details = model.get_output_details()
        if len(out_details_yaml) == len(output_details): return options_dict

        for out_detail in output_details:
            new_detail = tfl_detail_to_dict(out_detail)
            out_details_yaml.append(new_detail)

    return options_dict

def update_postprocess(options_dict):
    '''
    This doesn't handle all the variations in OD head's format, but it at least covers DetectionBoxSL2BoxLS
    '''
    postprocess = options_dict.get('postprocess', None)
    if postprocess is not None:
        formatter = postprocess.get('formatter', None)
        if formatter is not None:
            src_indices = formatter.get('src_indices', None)
            dst_indices = formatter.get('dst_indices', None)
            
            if src_indices is None and dst_indices is not None:
                formatter['src_indices'] = dst_indices.reverse()

            elif src_indices is not None and dst_indices is None:
                tmp = copy.copy(src_indices)
                tmp.reverse()
                formatter['dst_indices'] = tmp


    return options_dict



def get_model_path(options_dict, modeldir):
    modelpath = os.path.join(modeldir, options_dict['session']['model_path'])
    return modelpath

def open_model(modelpath):
    ext = os.path.splitext(modelpath)[-1]
    if ext == '.onnx':
        import onnxruntime as ort
        sess_options = ort.SessionOptions()
        model = ort.InferenceSession(
                    modelpath,
                    providers=['CPUExecutionProvider'], 
                    provider_options=[{}],
                    sess_options=sess_options,
                )
    #    
    elif ext == '.tflite':
        import tflite_runtime.interpreter as tfl
        model = tfl.Interpreter(modelpath)
    #
    return model, ext
        
def generate_dummy_dataset_yaml(num_classes, options_dict, modeldir, tasktype):
    postprocess = options_dict.get('postprocess')
    if postprocess is None: 
        postprocess = {}
        options_dict['postproces'] = postprocess

    label_pred = postprocess.get('label_offset_pred', None)
    if label_pred is not None:
        if tasktype == 'detection':
            num_defined_classes = len(label_pred.keys())
            print(num_defined_classes)
            num_classes = max(num_defined_classes, num_classes)
    else: 
        if tasktype == 'detection':
            label_pred = {}
            postprocess['label_offset_pred'] = label_pred
            num_defined_classes= 0
        elif tasktype == 'classification':
            label_pred = -1

    dataset_info = {
        'info': {},
        'categories': []
    }

    for i in range(num_classes):

        dataset_info['categories'].append({
                'id': i,
                'name': f'Class {i}' 
            })
        # if i > len(label_pred)+1: 
        if tasktype == 'detection':
            label_pred[i] = i

    yaml.dump(dataset_info, open(os.path.join(modeldir, 'dataset.yaml'), 'w'))
    options_dict['metric'] = {}
    options_dict['metric']['label_offset_pred'] = copy.copy(label_pred)

    return options_dict

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('modeldir', type=str, default=None)
    parser.add_argument('-n', '--num-classes', type=int, default=None)
    args = parser.parse_args()

    param_path = os.path.join(args.modeldir, 'param.yaml')


    with open(param_path, 'r') as param_file:
        param_yaml = yaml.safe_load(param_file)
        yaml.dump(param_yaml, open(os.path.join(args.modeldir, 'param-original.yaml'), 'w'))

        tasktype = param_yaml['task_type']

        if param_yaml.get('session', None) is None:
            param_yaml['session'] = {}

        param_yaml = update_preprocess(param_yaml)
        model, ext = open_model(get_model_path(param_yaml, args.modeldir))
        param_yaml = update_input_details(param_yaml, model, ext)
        param_yaml = update_output_details(param_yaml, model, ext)
        param_yaml = update_postprocess(param_yaml)
        if args.num_classes is not None and args.num_classes>0:
            param_yaml = generate_dummy_dataset_yaml(args.num_classes, param_yaml, args.modeldir, tasktype)

        yaml.dump(param_yaml, open(os.path.join(args.modeldir, 'param-updated.yaml'), 'w'))