text_plain_output_1.png
from monai import data

from monai import transforms

import json

import numpy as np

import os

import tempfile

import torch

directory = os.environ.get('MONAI_DATA_DIRECTORY')
root_dir = tempfile.mkdtemp() if directory is None else directory

class AverageMeter(object):
    def __init__(self):
        self.reset()
    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = np.where(self.count > 0, self.sum / self.count, self.sum)
def datafold_read(datalist, basedir, fold=0, key='training'):
    with open(datalist) as f:
        json_data = json.load(f)
    json_data = json_data[key]
    for d in json_data:
        for k, v in d.items():
            if isinstance(d[k], list):
                d[k] = [os.path.join(basedir, iv) for iv in d[k]]
            elif isinstance(d[k], str):
                d[k] = os.path.join(basedir, d[k]) if len(d[k]) > 0 else d[k]
    tr = []
    val = []
    for d in json_data:
        if 'fold' in d and d['fold'] == fold:
            val.append(d)
        else:
            tr.append(d)
    return (tr, val)
def save_checkpoint(model, epoch, filename='model.pt', best_acc=0, dir_add=root_dir):
    state_dict = model.state_dict()
    save_dict = {'epoch': epoch, 'best_acc': best_acc, 'state_dict': state_dict}
    filename = os.path.join(dir_add, filename)
    torch.save(save_dict, filename)

def get_loader(batch_size, data_dir, json_list, fold, roi):
    data_dir = data_dir
    datalist_json = json_list
    train_files, validation_files = datafold_read(datalist=datalist_json, basedir=data_dir, fold=fold)
    train_transform = transforms.Compose([transforms.LoadImaged(keys=['image', 'label']), transforms.ConvertToMultiChannelBasedOnBratsClassesd(keys='label'), transforms.CropForegroundd(keys=['image', 'label'], source_key='image', k_divisible=[roi[0], roi[1], roi[2]]), transforms.RandSpatialCropd(keys=['image', 'label'], roi_size=[roi[0], roi[1], roi[2]], random_size=False), transforms.RandFlipd(keys=['image', 'label'], prob=0.5, spatial_axis=0), transforms.RandFlipd(keys=['image', 'label'], prob=0.5, spatial_axis=1), transforms.RandFlipd(keys=['image', 'label'], prob=0.5, spatial_axis=2), transforms.NormalizeIntensityd(keys='image', nonzero=True, channel_wise=True), transforms.RandScaleIntensityd(keys='image', factors=0.1, prob=1.0), transforms.RandShiftIntensityd(keys='image', offsets=0.1, prob=1.0)])
    val_transform = transforms.Compose([transforms.LoadImaged(keys=['image', 'label']), transforms.ConvertToMultiChannelBasedOnBratsClassesd(keys='label'), transforms.NormalizeIntensityd(keys='image', nonzero=True, channel_wise=True)])
    train_ds = data.Dataset(data=train_files, transform=train_transform)
    train_loader = data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_ds = data.Dataset(data=validation_files, transform=val_transform)
    val_loader = data.DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
    return (train_loader, val_loader)

print(train_loader, val_loader)