#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import os
from Initial_Data_Program.Initial_Data_Gen_configs_base_datasets_my_dataset import generate_configs_base_datasets_my_dataset_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_utils_class_names import generate_mmseg_utils_class_names_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_datasets_my_dataset import generate_mmseg_datasets_my_dataset_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_datasets_init_ import generate_mmseg_datasets_init_file

if __name__ == '__main__':
    ########### 1.1.定义各数据集相关参数 ###########
    # 可以定义多个数据集文件名 和 对应类名
    dataset_file_names = ["my_dataset_model"] # =['my_dataset', 'my_dataset_2'] # =["my_dataset"]
    dataset_class_names = ["MyDataset_model"] # =['MyDataset', 'MyDataset2'] # =["MyDataset"]
    dataset_file_name='my_dataset_model' # 数据集 文件名.py   TODO
    dataset_class_name='MyDataset_model' # 数据集 类名称   TODO
    data_root='/home/audience/Desktop/Seg_data/Data' # 数据根目录
    img_scale=(1920, 1080) # 图片大小
    # 训练、验证、测试集所在文件夹
    train_img_path='A_Ori'
    train_seg_map_path='A_Label_GT_label_fold'
    val_img_path='A_Ori'
    val_seg_map_path='A_Label_GT_label_fold'
    test_img_path='A_Ori'
    test_seg_map_path='A_Label_GT_label_fold'
    
    ########### 1.2.定义Label图片相关参数 ###########
    classes = ['肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉']
    palette = [[255,91,0],[255,234,0],[85, 111, 181],[181, 227, 14],[72, 0, 255],[0, 155, 33],[255,0,255],[29, 32, 136],[160, 15, 95],[0,160,233],[52,184,178],[90,120,41],[255,0,0],[177,0,0],[167,24,233],[112,113,150],[0,255,0],[255,255,255],[0,255,255],[138,251,213],[136,162,196],[197,83,181],[202,202,200],[113,102,140],[66,115,82],[240,16,116],[155,132,0],[155,62,0],[146,175,236],[255,172,159],[245,161,0],[134,124,118],[0,157,142],[181,85,105],[42,8,66]]
    # 这里的classes一定是经过“Initial_Gen_mmseg_datasets_my_dataset.py”处理的
    classes_all = [
        ['背景','肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉'],
    ]
    palette_all = [
        [[0,0,0],[255,91,0],[255,234,0],[85, 111, 181],[181, 227, 14],[72, 0, 255],[0, 155, 33],[255,0,255],[29, 32, 136],[160, 15, 95],[0,160,233],[52,184,178],[90,120,41],[255,0,0],[177,0,0],[167,24,233],[112,113,150],[0,255,0],[255,255,255],[0,255,255],[138,251,213],[136,162,196],[197,83,181],[202,202,200],[113,102,140],[66,115,82],[240,16,116],[155,132,0],[155,62,0],[146,175,236],[255,172,159],[245,161,0],[134,124,118],[0,157,142],[181,85,105],[42,8,66]],
    ]
    # 一般不太会变的参数
    img_suffix = ".png" 
    seg_map_suffix = "_gtFine_labelTrainIds.png"
    reduce_zero_label = False # 在第 0 类为 background 类别的数据集上，如果您最终是需要将背景和您的其余类别分开时，是不需要使用reduce_zero_label的 【reduce_zero_label = False】

    ########### 1.3.定义训练相关参数 ###########
    # 一般不太会变的参数
    crop_size=(512, 512) # 分割大小
    train_batch_size=4 # 训练batch
    train_num_workers=4 # 训练并行运行数量
    val_and_test_batch_size=1 # 验证集和测试集batch
    val_and_test_num_workers=4 # 验证集和测试集并行运行数量
    
    ########### 2.文件存储位置 ###########
    output_configs_base_datasets_my_dataset=f'./configs/_base_/datasets/{dataset_file_name}.py'
    output_mmseg_datasets_dataset_file_name = os.path.join(f'./mmseg/datasets/{dataset_file_name}.py')
    output_mmseg_datasets_init = os.path.join('./mmseg/datasets/__init__.py')
    output_mmseg_utils_class_names = f'./mmseg/utils/class_names.py'

    ########### 3.运行程序生成配置文件 ###########
    success = generate_configs_base_datasets_my_dataset_file(output_file=output_configs_base_datasets_my_dataset, dataset_class_name=dataset_class_name , data_root=data_root, img_scale=img_scale, crop_size=crop_size, train_batch_size=train_batch_size, train_num_workers=train_num_workers, val_and_test_batch_size=val_and_test_batch_size, val_and_test_num_workers=val_and_test_num_workers, train_img_path=train_img_path, train_seg_map_path=train_seg_map_path, val_img_path=val_img_path, val_seg_map_path=val_seg_map_path, test_img_path=test_img_path, test_seg_map_path=test_seg_map_path)
    success, classes, palette = generate_mmseg_datasets_my_dataset_file(output_file=output_mmseg_datasets_dataset_file_name, dataset_class_name=dataset_class_name, classes=classes, palette=palette, img_suffix=img_suffix, seg_map_suffix=seg_map_suffix, reduce_zero_label=reduce_zero_label)
    # 需要用到上一步的classes和palette
    success = generate_mmseg_datasets_init_file(output_file=output_mmseg_datasets_init, dataset_file_names=dataset_file_names, dataset_class_names=dataset_class_names)
    success = generate_mmseg_utils_class_names_file(output_file=output_mmseg_utils_class_names, dataset_file_names=dataset_file_names, classes_all=classes_all, palette_all=palette_all)

