web3-202226701039/Main(2)(2).py

import os
import numpy as np
import time
import sys

from ChexnetTrainer import ChexnetTrainer  # 从ChexnetTrainer模块中导入训练相关的类


# --------------------------------------------------------------------------------

# 主函数，负责启动不同的功能（训练、测试或运行演示）
def main():
    #runDemo()  # 运行演示模式
    # runTest()  # 测试模式（注释掉，可以通过解除注释来运行）
    runTrain()  # 训练模式（注释掉，可以通过解除注释来运行）


# --------------------------------------------------------------------------------

# 训练函数，定义训练所需的参数并启动模型训练
def runTrain():
    DENSENET121 = 'DENSE-NET-121'  # 定义DenseNet121模型名称
    DENSENET169 = 'DENSE-NET-169'  # 定义DenseNet169模型名称
    DENSENET201 = 'DENSE-NET-201'  # 定义DenseNet201模型名称
    Resnet50='RESNET-50'  # 定义Resnet50模型名称

    # 获取当前的时间戳，作为训练过程的标记
    timestampTime = time.strftime("%H%M%S")
    timestampDate = time.strftime("%d%m%Y")
    timestampLaunch = timestampDate + '-' + timestampTime
    print("Launching " + timestampLaunch)

    # 图像数据所在的路径
    pathDirData = 'chest xray14'

    # 训练、验证和测试数据集文件路径
    # 每个文件中包含图像路径及其对应的标签
    pathFileTrain = './dataset/train_2.txt'
    pathFileVal = './dataset/valid_2.txt'
    pathFileTest = './dataset/test_2.txt'

    # 神经网络参数：模型架构、是否加载预训练模型、分类的类别数量
    nnArchitecture = DENSENET121
    nnIsTrained = True  # 使用预训练的权重
    nnClassCount = 14  # 数据集包含14个分类

    # 训练参数：批量大小和最大迭代次数（epochs）
    trBatchSize = 16
    trMaxEpoch = 48

    # 图像预处理相关参数：图像缩放的大小和裁剪后的大小
    imgtransResize = 256
    imgtransCrop = 224

    # 保存模型的路径，包含时间戳

    pathModel = 'm-' + timestampLaunch + '.pth.tar'

    print('Training NN architecture = ', nnArchitecture)
    ChexnetTrainer.train(pathDirData, pathFileTrain, pathFileVal,
                         nnArchitecture, nnIsTrained, nnClassCount, trBatchSize,
                         trMaxEpoch, imgtransResize, imgtransCrop,
                         timestampLaunch, None)
    print('Testing the trained model')
    ChexnetTrainer.test(pathDirData, pathFileTest, pathModel, nnArchitecture,
                        nnClassCount, nnIsTrained, trBatchSize, imgtransResize,
                        imgtransCrop, timestampLaunch)

# --------------------------------------------------------------------------------

# 测试函数，加载预训练模型并在测试数据集上进行测试
def runTest():
    pathDirData = '/chest xray14'  # 数据路径
    pathFileTest = './dataset/test.txt'  # 测试集路径
    nnArchitecture = 'DENSE-NET-121'  # 使用DenseNet121架构
    nnIsTrained = True  # 使用预训练模型
    nnClassCount = 14  # 分类数
    trBatchSize = 4  # 批量大小
    imgtransResize = 256  # 图像缩放大小
    imgtransCrop = 224  # 图像裁剪大小

    # 预训练模型路径
    pathModel = 'm-06102024-235412BCELoss()delete.pth.tar'

    timestampLaunch = ''  # 时间戳

    # 调用测试函数，使用上述参数
    ChexnetTrainer.test(pathDirData, pathFileTest, pathModel, nnArchitecture,
                        nnClassCount, nnIsTrained,
                        trBatchSize, imgtransResize, imgtransCrop,
                        timestampLaunch)


# --------------------------------------------------------------------------------

# 演示函数，展示模型在测试集上的推理过程
def runDemo():
    pathDirData = '/media/sunjc0306/未命名/CODE/chest xray14'  # 数据路径
    pathFileTest = './dataset/test.txt'  # 测试集路径
    nnArchitecture = 'DENSE-NET-121'  # 使用DenseNet121架构
    nnIsTrained = True  # 使用预训练模型
    nnClassCount = 14  # 分类数
    trBatchSize = 4  # 批量大小
    imgtransResize = 256  # 图像缩放大小
    imgtransCrop = 224  # 图像裁剪大小

    pathModel = 'm-06102024-235412BCELoss()delete.pth.tar'  # 预训练模型路径

    # 定义分类名称
    CLASS_NAMES = ['Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration',
                   'Mass', 'Nodule', 'Pneumonia',
                   'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema',
                   'Fibrosis', 'Pleural_Thickening', 'Hernia']

    import torch
    import torch.backends.cudnn as cudnn
    import torchvision.transforms as transforms
    from torch.utils.data import DataLoader
    from tqdm import tqdm  # 用于进度条显示
    from DensenetModels import DenseNet121  # 导入模型
    from DensenetModels import DenseNet169
    from DensenetModels import DenseNet201
    from DatasetGenerator import DatasetGenerator  # 导入数据集生成器
    cudnn.benchmark = True  # 加速模型的推理过程

    # 设置网络架构并加载模型
    if nnArchitecture == 'DENSE-NET-121':
        model = DenseNet121(nnClassCount, nnIsTrained).cuda()
    elif nnArchitecture == 'DENSE-NET-169':
        model = DenseNet169(nnClassCount, nnIsTrained).cuda()
    elif nnArchitecture == 'DENSE-NET-201':
        model = DenseNet201(nnClassCount, nnIsTrained).cuda()

    model = model.cuda()

    # 加载预训练的模型权重
    modelCheckpoint = torch.load(pathModel)
    model.load_state_dict(modelCheckpoint['state_dict'])

    # 图像变换：先缩放，再进行裁剪，最后进行归一化
    normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])

    # 数据集变换和数据加载器
    transformList = []
    transformList.append(transforms.Resize(imgtransResize))
    transformList.append(transforms.TenCrop(imgtransCrop))  # 对图像进行十裁剪
    transformList.append(transforms.Lambda(lambda crops: torch.stack(
        [transforms.ToTensor()(crop) for crop in crops])))  # 转换为张量
    transformList.append(transforms.Lambda(
        lambda crops: torch.stack([normalize(crop) for crop in crops])))  # 归一化
    transformSequence = transforms.Compose(transformList)

    # 构建测试集数据集生成器
    datasetTest = DatasetGenerator(pathImageDirectory=pathDirData,
                                   pathDatasetFile=pathFileTest,
                                   transform=transformSequence)

    model.eval()  # 设置模型为评估模式
    i = 0  # 初始索引

    # 定义演示函数，展示模型在某个样本上的推理结果
    def demo(i):
        (input, target) = datasetTest[i]  # 获取输入和目标
        n_crops, c, h, w = input.size()  # 获取输入的尺寸

        # 将输入放入模型进行推理
        varInput = torch.autograd.Variable(input.view(-1, c, h, w).cuda(),
                                           volatile=True)
        with torch.no_grad():
            out = model(varInput)
            outMean = out.view(1, n_crops, -1).mean(1)  # 对裁剪后的多个输出取平均
        print('-------------------------------')
        pd = torch.sigmoid(outMean).ge_(0.5).long().detach().cpu().numpy()[
            0].tolist()  # 预测结果
        gt = target.long().detach().cpu().numpy().tolist()  # 真实标签
        print(f"PD_{i}", pd)  # 输出预测结果
        print(f"GT_{i}", gt)  # 输出真实标签
        return pd, gt

    # 逐个演示测试集中的样本
    for i in range(len(datasetTest)):
        pd, gt = demo(i)  # 调用演示函数
        if type(gt) == list:
            pd_name = []
            gt_name = []
            for i in range(len(CLASS_NAMES)):
                if pd[i] == 1:
                    pd_name.append(CLASS_NAMES[i])  # 将预测为1的标签存入pd_name
                if gt[i] == 1:
                    gt_name.append(CLASS_NAMES[i])  # 将真实为1的标签存入gt_name
            print(datasetTest.listImagePaths[i])  # 打印图像路径
            print(pd_name)  # 打印预测的疾病名称
            print(gt_name)  # 打印真实的疾病名称


if __name__ == '__main__':
    main()  # 启动主函数