Algorithm-Repo/SourceCode/MachineLearningUtils/TestTools/AIModelMetricsUtils/predict_file_generate.py

"""
将预测结果，装换成所需要的pred格式文件，
为了调用方便，pred格式文件和gt格式文文件一致
分为分类，检测，语义分割，三种不同的格式
后续添加实例分割的格式，直接添加函数即可
"""

"""
将所有c#返回的DetectedObject格式，根据模型类型（分类，检测，语义分割）转换成所需的格式，如下:

分类：
[{'Label': 0, 'Confidence': 1.0}]
[]

检测：
[{'Label': 1, 'Confidence': 1.0, 'BoundingBox'：[2, 3, 4, 5]}]
[{'Label': 1, 'Confidence': 1.0, 'BoundingBox'：[2, 3, 4, 5]},{'Label': 2, 'Confidence': 1.0, 'BoundingBox'：[2, 3, 4, 5]}]
[{'Label': 0, 'Confidence': 1.0, 'BoundingBox'：[0, 0, 0, 0]}]
[]

语义分割：
[{'Label': 0, 'Confidence': 1.0, 'Image_size':[w, h], 'Contour': []}]
[{'Label': 1, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours'：[ [ [1,2],[3,4] ] , [ [5,6],[7,8] ]  ]}]
[{'Label': 1, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours'：[ [ [1,2],[3,4] ]  ]}]
[{'Label': 1, 'Confidence': 1.0, ''Image_size':[w, h], Contours'：[ [ [1,2],[3,4] ]  ]},{'Label': 2, 'Confidence': 1.0, 'Contours'：[ [ [1,2],[3,4] ]  ]}]
[]

"""
import json


def predfilegenerate_classification(idetectedobject, image_size, class_id_map):
    """
    设置所需的gt格式,返回需要的json格式
    适用于: 分类模型

    :param idetectedobject: 模型调用结果
    :param image_size: 图像尺寸
    :param class_id_map: 用于label改变的dict
    :return: 返回分类所需的文件格式
    """
    output_info_list = []
    length = idetectedobject.Length
    if length == 0:
        output_info = {'Label': 0,
                       'Confidence': 1.0}
        output_info_list.append(output_info)
    if length == 1:
        for j in range(length):
            output_info = {'Label': 0,
                           'Confidence': 1.0}
            label = class_id_map[idetectedobject[j].Label]
            confidence = idetectedobject[j].Confidence
            output_info['Label'] = label
            output_info['Confidence'] = confidence
            output_info_list.append(output_info)
    # 保证输出的output_info_list
    if len(output_info_list) == 0:
        print('label_info为：{} 的图像，无法匹配任何所需的imageresults和rois的title，该图像作废'.format(idetectedobject))
        return json.dumps([], ensure_ascii=False)
    else:
        return json.dumps(output_info_list, ensure_ascii=False)


def predfilegenerate_object_detection(idetectedobject, image_size, class_id_map):
    """
    设置所需的gt格式,返回需要的json格式
    适用于: 检测模型

    :param idetectedobject:
    :param image_size:
    :param class_id_map:
    :return: 返回检测所需的文件格式
    """
    output_info_list = []
    length = idetectedobject.Length
    if length == 0:
        output_info = {'Label': 0,
                       'Confidence': 1.0,
                       'BoundingBox': [0, 0, 0, 0]}
        output_info_list.append(output_info)
    else:
        for j in range(length):
            output_info = {'Label': 0,
                           'Confidence': 1.0,
                           'BoundingBox': [0, 0, 0, 0]}
            label = class_id_map[idetectedobject[j].Label]
            confidence = idetectedobject[j].Confidence
            rect_BoundingBox = idetectedobject[j].BoundingBox
            rect_top = rect_BoundingBox.Top
            rect_left = rect_BoundingBox.Left
            rect_right = rect_BoundingBox.Right
            rect_bottom = rect_BoundingBox.Bottom

            output_info['Label'] = label
            output_info['BoundingBox'] = [rect_left, rect_top, rect_right, rect_bottom]
            output_info['Confidence'] = confidence
            output_info_list.append(output_info)

    # 保证输出的output_info_list
    if len(output_info_list) == 0:
        print('label_info为：{} 的图像，无法匹配任何所需的imageresults和rois的title，该图像作废'.format(idetectedobject))
        return json.dumps([], ensure_ascii=False)
    else:
        return json.dumps(output_info_list, ensure_ascii=False)


def predfilegenerate_semantic_segmentation(idetectedobject, image_size, class_id_map):
    """
    设置所需的gt格式,返回需要的json格式
    适用于: 语义分割

    :param label_info:
    :param needed_imageresults_dict:
    :param needed_rois_dict:
    :return:

    """
    output_info_list = []
    length = idetectedobject.Length
    if length == 0:
        output_info = {'Label': 0,
                       'Image_size': image_size,
                       'Confidence': 1.0,
                       'Contours': []}
        output_info_list.append(output_info)

    for j in range(length):
        output_info = {'Label': 0,
                       'Image_size': image_size,
                       'Confidence': 1.0,
                       'Contours': []}
        label = class_id_map[idetectedobject[j].Label]
        confidence = idetectedobject[j].Confidence
        rect_BoundingBox = idetectedobject[j].BoundingBox
        contours = idetectedobject[j].Contours
        for z in range(contours.Length):
            each_contour = []
            for zi in range(contours[z].Length):
                each_contour.append([contours[z][zi].X, contours[z][zi].Y])
            output_info['Contours'].append(each_contour)
        output_info['Label'] = label
        output_info['Confidence'] = confidence
        output_info_list.append(output_info)

    # 保证输出的output_info_list
    if len(output_info_list) == 0:
        print('label_info为：{} 的图像，无法匹配任何所需的imageresults和rois的title，该图像作废'.format(idetectedobject))
        return json.dumps([], ensure_ascii=False)
    else:
        return json.dumps(output_info_list, ensure_ascii=False)


def predfilegenerate_semantic_segmentation_myocardial(idetectedobject, image_size, class_id_map):
    """
    设置所需的gt格式,返回需要的json格式
    适用于: 心肌分割
    原因是：心肌分割只输出一个轮廓，并且label为0，自动+1，label设为1

    :param label_info:
    :param needed_imageresults_dict:
    :param needed_rois_dict:
    :return:

    """
    output_info_list = []
    length = idetectedobject.Length
    if length == 0:
        output_info = {'Label': 0,
                       'Image_size': image_size,
                       'Confidence': 1.0,
                       'Contours': []}
        output_info_list.append(output_info)

    for j in range(length):
        output_info = {'Label': 0,
                       'Image_size': image_size,
                       'Confidence': 1.0,
                       'Contours': []}
        label = class_id_map[idetectedobject[j].Label + 1]
        confidence = idetectedobject[j].Confidence
        rect_BoundingBox = idetectedobject[j].BoundingBox
        contours = idetectedobject[j].Contours
        for z in range(contours.Length):
            each_contour = []
            for zi in range(contours[z].Length):
                each_contour.append([contours[z][zi].X, contours[z][zi].Y])
            output_info['Contours'].append(each_contour)
        output_info['Label'] = label
        output_info['Confidence'] = confidence
        output_info_list.append(output_info)

    # 保证输出的output_info_list
    if len(output_info_list) == 0:
        print('label_info为：{} 的图像，无法匹配任何所需的imageresults和rois的title，该图像作废'.format(idetectedobject))
        return json.dumps([], ensure_ascii=False)
    else:
        return json.dumps(output_info_list, ensure_ascii=False)