Update README.md

Author: Stefan Knegt

README.md

@@ -37,8 +37,22 @@ for epoch in range(epochs):
 ## Train on LIDC Dataset
 One of the datasets used in the original paper is the [LIDC dataset](https://wiki.cancerimagingarchive.net). I've preprocessed this data and stored them in 5 .pickle files which you can [download here](https://drive.google.com/drive/folders/1xKfKCQo8qa6SAr3u7qWNtQjIphIrvmd5?usp=sharing). After downloading the files you can load the data as follows:
 ```
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.sampler import SubsetRandomSampler
 from load_LIDC_data import LIDC_IDRI
+
 dataset = LIDC_IDRI(dataset_location = 'insert_path_here')
+dataset_size = len(dataset)
+indices = list(range(dataset_size))
+split = int(np.floor(test_split * dataset_size))
+np.random.shuffle(indices)
+train_indices, test_indices = indices[split:], indices[:split]
+train_sampler = SubsetRandomSampler(train_indices)
+test_sampler = SubsetRandomSampler(test_indices)
+train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler)
+test_loader = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler)
+print("Number of training/test patches:", (len(train_indices),len(test_indices)))
 ```
 Combining this with the training code snippet above, you can start training your own Probabilistic UNet.