Algorithm-Repo/SourceCode/MachineLearningUtils/TestTools/AITestDemo/PythonPart/CsharpTestTxtGenerate.py

# D:/workplace/python
# -*- coding: utf-8 -*-
# @File  :CsharpTestTxtGenerate.py
# @Author:Guido    LuXiaohao
# @Date  :2021/11/8
# @Software:PyCharm

"""
C# AITestDemo 测试所需txt，输入格式是AI训练平台上的json格式
"""

import os
import json
import cv2
import shutil

import numpy as np

def ClassTxtGenerate(source_path, output_path, EnumDetectedObjectType, output_label_dict):
    # 包含图片和对应label的txt文件路径（其中txt格式为AI平台训练的json格式）
    for i in os.listdir(source_path):
        if ".txt" in i:
            source_name = i[:-4]
            # print(source_name)
            txt_path = os.path.join(source_path, source_name + ".txt")
            # 防止覆写
            if os.path.exists(os.path.join(output_path, source_name + ".txt")):
                print("File already exist!")
                return
            output_txt_file = open(os.path.join(output_path, source_name + ".txt"), "a", encoding="utf-8")
            output_info_list = []
            with open(txt_path, encoding="utf-8") as f:
                output_info = {'EnumDetectedObjectType': None,
                               'Label': 0,
                               'Confidence': 1.0,
                               'BoundingBox': None,
                               'Contour': []}
                image = cv2.imdecode(np.fromfile(os.path.join(source_path, source_name + ".jpg"), dtype=np.uint8), -1)
                image_height, image_width = image.shape[0], image.shape[1]
                print(image.shape[0], image.shape[1])
                BoundingBox = {'Left': 0, 'Top': 0, 'Right': image_width, 'Bottom': image_height, 'Width': image_width, 'Height': image_height}

                label_info = json.loads(f.readline())
                # print(label_info)
                for result_index in range(len(label_info["ImageResults"])):
                    labelname = label_info["ImageResults"][result_index]["Conclusion"]["Title"]
                    # print(labelname)
                    if labelname not in output_label_dict.keys():
                        print(labelname)
                    else:
                        output_info['EnumDetectedObjectType'] = EnumDetectedObjectType
                        output_info['Label'] = output_label_dict[labelname]
                        output_info['BoundingBox'] = BoundingBox
                    output_info_list.append(output_info)
            shutil.copy(os.path.join(source_path, source_name + ".jpg"),
                        os.path.join(output_path, source_name + ".jpg"))
            output_txt_file.write('{}'.format(json.dumps(output_info_list, ensure_ascii=False)))
            output_txt_file.close()

def SemanticSegTxtGenerate(source_path, output_path, EnumDetectedObjectType,
                           crop_organ_label_dict, output_label_dict):
    # 包含图片和对应label的txt文件路径（其中txt格式为AI平台训练的json格式）
    for source_index, source in enumerate(os.listdir(source_path)):

        if ".txt" in source:
            source_name = source[:-4]
            txt_path = os.path.join(source_path, source_name + ".txt")
            # 防止覆写
            if os.path.exists(os.path.join(output_path, source_name + ".txt")):
                print("File already exist!")
                return
            output_txt_file = open(os.path.join(output_path, source_name + ".txt"), "a", encoding="utf-8")
            output_info_list = []
            with open(txt_path, encoding="utf-8") as f:
                image = cv2.imdecode(np.fromfile(os.path.join(source_path, source_name + ".jpg"), dtype=np.uint8), -1)
                image_height, image_width = image.shape[0], image.shape[1]
                label_info = json.loads(f.readline())

                max_roi_area = 0
                crop_output_info = {'EnumDetectedObjectType': None,
                                    'Label': -1,
                                    'Confidence': 1.0,
                                    'BoundingBox': {'Left': 0, 'Top': 0, 'Right': image_width, 'Bottom': image_height, 'Width': image_width, 'Height': image_height},
                                    'Contour': [{'x': 0, 'y': 0}, {'x': image_width, 'y': 0}, {'x': image_width, 'y': image_height}, {'x': 0, 'y': image_height}]}
                for roi_index, each_roi_label in enumerate(label_info["Rois"]):
                    output_info = {'EnumDetectedObjectType': None,
                                   'Label': 0,
                                   'Confidence': 1.0,
                                   'BoundingBox': None,
                                   'Contour': []}
                    labelname = each_roi_label["Conclusion"]["Title"]
                    output_info['EnumDetectedObjectType'] = EnumDetectedObjectType
                    if labelname in output_label_dict.keys():
                        roi_points = each_roi_label["Points"]
                        x, y = [], []
                        roi_contour = []
                        for point in roi_points:
                            x.append(int(point.split(",")[0]))
                            y.append(int(point.split(",")[1]))
                            roi_contour.append({'x': int(point.split(",")[0]),
                                                'y': int(point.split(",")[1])})
                        left, right = min(x), max(x)
                        top, bottom = min(y), max(y)
                        roi_boundingbox = {'Left': left, 'Top': top, 'Right': right, 'Bottom': bottom,
                                           'Width': right - left, 'Height': bottom - top}
                        # 赋值
                        output_info['Label'] = output_label_dict[labelname]
                        output_info['BoundingBox'] = roi_boundingbox
                        output_info['Contour'] = roi_contour
                        output_info_list.append(output_info)
                    elif labelname in crop_organ_label_dict.keys():
                        roi_points = each_roi_label["Points"]
                        x, y = [], []
                        roi_contour = []
                        for point in roi_points:
                            x.append(int(point.split(",")[0]))
                            y.append(int(point.split(",")[1]))
                            roi_contour.append({'x': int(point.split(",")[0]),
                                                'y': int(point.split(",")[1])})
                        left, right = min(x), max(x)
                        top, bottom = min(y), max(y)
                        roi_boundingbox = {'Left': left, 'Top': top, 'Right': right, 'Bottom': bottom,
                                           'Width': right - left, 'Height': bottom - top}
                        roi_area = (right - left) * (bottom - top)
                        if roi_area > max_roi_area:
                            crop_output_info['Label'] = crop_organ_label_dict[labelname]
                            crop_output_info['BoundingBox'] = roi_boundingbox
                            crop_output_info['Contour'] = roi_contour
                            max_roi_area = roi_area
                if crop_organ_label_dict != {}:
                    crop_output_info['EnumDetectedObjectType'] = EnumDetectedObjectType
                    # 裁切部位放第一位
                    output_info_list.insert(0, crop_output_info)

            shutil.copy(os.path.join(source_path, source_name + ".jpg"),
                        os.path.join(output_path, source_name + ".jpg"))
            output_txt_file.write('{}'.format(json.dumps(output_info_list, ensure_ascii=False)))
            output_txt_file.close()

if __name__ == '__main__':
    # #######分类用######
    # source_path = r"D:\Vinno_Guido\Code\文件\训练格式\TestImages"
    # output_path = r"D:\Vinno_Guido\Code\文件\训练格式\output"
    # EnumDetectedObjectType = "Scanpart"
    # output_label_dict = {"腹部B模式_vinno": 0,
    #                      "乳腺B模式_vinno": 1,
    #                      "超声背景图": 3, "产科B模式_vinno": 3,
    #                      "甲状腺B模式_vinno": 4}
    # ClassTxtGenerate(source_path, output_path, EnumDetectedObjectType, output_label_dict)

    # #######分割用(无需裁切部位)######
    # source_path = r"E:\Vinno\数据\AnnotatedLiverDatas_20210611\本地调试1"
    # output_path = r"C:\Users\VINNO\Desktop\Csharptxt"
    # EnumDetectedObjectType = r"Organ"
    # crop_organ_label_dict = {}
    # output_label_dict = {"肝": 1,
    #                      "胆囊胆道": 2,
    #                      "肾脏": 3,
    #                      "脾脏": 4}
    # SemanticSegTxtGenerate(source_path, output_path, EnumDetectedObjectType,
    #                        crop_organ_label_dict, output_label_dict)

    # #######分割用(需裁切部位)######
    # source_path = r"E:\Vinno\数据\AnnotatedLiverDatas_20210611\本地调试1"
    # output_path = r"C:\Users\VINNO\Desktop\Csharptxt"
    # EnumDetectedObjectType = r"Lesion"
    # crop_organ_label_dict = {"肝脏未见明显异常": -1,
    #                          "肝脏有局灶性疾病/区域": -1,
    #                          "脂肪肝声像图改变": -1,
    #                          "肝脏弥漫性病变声像图改变": -1,
    #                          "肝硬化声像图改变": -1,
    #                          "多囊肝声像图改变": -1}
    # output_label_dict = {"肝内强回声灶": 1,
    #                      "肝血管瘤声像图改变": 2,
    #                      "肝囊肿": 3,
    #                      "肝癌可能": 4}
    # SemanticSegTxtGenerate(source_path, output_path, EnumDetectedObjectType,
    #                        crop_organ_label_dict, output_label_dict)

    pass