Beispiel #1
0
def main():
    # Configs
    config = configs.config()

    # Reading the selected data
    DataCSVFrame = pd.read_csv("DataFrame.csv",
                               usecols=["Image_Index", "Finding_Labels"],
                               index_col=False)
    labelsSet = set(DataCSVFrame["Finding_Labels"].values)

    # Dictionary with the label as key and the index in the set as value
    labelsDict = {}
    # Dictionary that is the reverse of the one above, to change back from value to the corresponding label
    labelDictClassify = {}

    # Filling the dictionaries
    for index, label in enumerate(labelsSet):
        labelsDict[label] = index
        labelDictClassify[index] = label

    # Path where all the images are stored
    imgPath = config.getImagePath()
    # Creating the dataset
    xrayDataset = DC.XRayDataset(DataCSVFrame, imgPath, labelsDict)

    #    # Getting the first image from the dataset
    #    imgs, labs = xrayDataset.__getitem__([8307])
    #    print(len(imgs))
    #    print(imgs)

    # Get the device (cpu/gpu) to run the model on
    device = DU.getDevice()

    # Gets the ranges of training and test data
    training, testing = DU.splitTrainTest(xrayDataset, config)

    # Get the train and validation sets
    trainSets, valSets = DU.trainValSets(training, config)

    # Initialize the criterion, optimizer and model
    criterion, optimizer, model = NM.modelInit(device)

    # Get the batchsize
    batchsize = config.getBatchSize()

    # Train the model
    trainedModel = NM.trainNetwork(device, xrayDataset, trainSets, valSets,
                                   config, model, criterion, optimizer,
                                   batchsize)

    # Save the model to be used for testing
    NM.save_model(trainedModel, config.getModelName())
Beispiel #2
0
def main():
    # Configs
    config = configs.config()

    # Reading the selected data
    DataCSVFrame = pd.read_csv("DataFrame.csv",
                               usecols=["Image_Index", "Finding_Labels"],
                               index_col=False)

    # Get the labels
    labelsSet = set(DataCSVFrame["Finding_Labels"].values)

    # Dictionary with the label as key and the index in the set as value
    labelsDict = {}
    # Dictionary that is the reverse of the one above, to change back from value to the corresponding label
    labelDictClassify = {}

    # Filling the dictionaries
    for index, label in enumerate(labelsSet):
        labelsDict[label] = index
        labelDictClassify[index] = label

    # Path where all the images are stored
    imgPath = config.getImagePath()
    # Creating the dataset
    xrayDataset = DC.XRayDataset(DataCSVFrame, imgPath, labelsDict)

    #print(xrayDataset.xrayClassFrame)

    # Get the device (cpu/gpu) to run the model on
    device = DU.getDevice()

    # Gets the ranges of training and test data
    training, testing = DU.splitTrainTest(xrayDataset, config)

    # Initialize the model
    model = models.alexnet(pretrained=False, num_classes=4)
    model.to(device)
    if device == 'cuda':
        model = torch.nn.DataParallel(model)
        cudnn.benchmark = True

    #print(model)

    # Load the trained model
    cwd = os.path.dirname(os.path.realpath(__file__))
    model.load_state_dict(
        torch.load("%s%s%s.pth" % (cwd, os.sep, config.getModelName())))
    model.eval()

    # Testing the model
    wrongLabels, labelsCM, predsCM = NM.testing(xrayDataset, testing, model,
                                                device, labelDictClassify)

    # Confusion Matrix
    CMPlot.plot_confusion_matrix(labelsCM,
                                 predsCM,
                                 list(labelsSet),
                                 normalize=False,
                                 title="Confusion Matrix")

    # Show the matrix
    plt.show()