This repository contains PyTorch code for cross-validation and inference with neural networks for kidney segmentation on UK Biobank neck-to-knee body MRI, as described in:

[_"Kidney segmentation in neck-to-knee body MRI of 40,000 UK Biobank participants"_](https://arxiv.org/abs/2006.06996) [1]

Contents:

- 2.5D U-Net architecture with residual connections

- Infrastructure for training and *cross-validation*

- Code for *quality_controls* based on numerical metrics

- Trained snapshot for parenchymal kidney tissue can be found at *TODO*

3. Run *crossValidate.py* to train and evaluate a model, with the results stored to the directory "networks"

To re-train a network for inference using all data, the cross-validation split can simply be set to contain all images in one split set.

split_path = "../splits/" + split_name + "/"

id_list = split_path + "id_list.txt"

def main(argv):

#

#####

createFolders(output_path, overwrite=True)

(slices_seg, shape_seg) = formatSeg(path_seg + file_s, shape_img)

if not np.array_equal(shape_img, shape_seg):

sys.exit()

# For each axial slice, save outputs

import dataLoading

def main(argv):

path_network_out = "../networks/kidney_64_8fold/"

path_stations_img = "/media/taro/DATA/Taro/UKBiobank/segmentations/kidney/combined_128/signals/NRRD/"

path_stations_gt = "/media/taro/DATA/Taro/UKBiobank/segmentations/kidney/combined_128/segmentations/NRRD_fixedHeaders/"

# Set to None for conventional cross-validation

path_train_ids_add = None

I = 80000 # Training iterations

save_step = 5000 # Iterations between checkpoint saving

class_count = 2 # Number of labels, including background

class_weights = torch.FloatTensor([1, 1]) # Background, L1, L2...

do_train = True

do_predict = True

os.makedirs(path_out_k)

os.makedirs(path_checkpoints)

time = train.train(net, loader_train, I, path_checkpoints, save_step, class_weights, I_reduce_lr)

with open(path_out_k + "training_time.txt", "w") as f: f.write("{}".format(time))

if do_predict:

evaluate.writeSubsetTrainingCurve(path_out_k)

# Get chosen volumes

subject_ids = []

print("Loading data for {} subjects".format(len(subject_ids)))

stations = []

# Get training samples

files = [f for f in os.listdir(input_path) if os.path.isfile(os.path.join(input_path, f))]

plt.close()

time_start = time.time()

# Fuse and store reference segmentation

for i in range(N):

# Find checkpoints

checkpoint_files = [f for f in os.listdir(path_checkpoints) if os.path.isfile(os.path.join(path_checkpoints, f))]

print(" Evaluating snapshot {}...".format(i))

checkpoint_i = path_checkpoints + checkpoint_files[i]

iteration = checkpoint_files[i].split("_")[1].split(".")[0]

evaluateAgreement(path_out, iteration, val_subjects)

voxel_dim = np.array((space_dir[0][0], space_dir[1][1], space_dir[2][2]))

voxel_scale = np.prod(voxel_dim) / (10*10*10)

# Get positives, true positives, false positives

p = np.count_nonzero(gt)

tp = np.count_nonzero(np.multiply(gt, out))

f.write("{},{},{},{},{},{}\n".format(val_subjects[i],dice,p,tp,fp,voxel_scale))

# Load network weights

checkpoint = torch.load(path_checkpoint, map_location={"cuda" : "cpu"})

print("Subject {}".format(val_subjects[i]))

if not os.path.exists(path_out + "{}_img.nrrd".format(val_subjects[i])):

nrrd.write(path_out + "{}_img.nrrd".format(val_subjects[i]), img, header, compression_level=1)

nrrd.write(path_out + "{}_out.nrrd".format(val_subjects[i]), out, header, compression_level=1)

header["sizes"] = gt.shape

nrrd.write(path_out + "{}_gt.nrrd".format(subject_id), gt, header, compression_level=1)

c_use_gpu = True # If yes, use numba for gpu access, otherwise use scipy on cpu

if not os.path.exists(input_path_zip):

print("Could not find input file {}".format(input_path_zip))

headers = []

# Only use abdominal imaging stations

#

voxel_data_w[i] = np.flip(voxel_data_w[i], 2)