Dataset for Prostate Cancer

README.md

@@ -1,3 +1,38 @@
+## Branch info
+
+Dataset for [Prostate Cancer Dataset](https://promise12.grand-challenge.org)
+
+### Usage
+
+```python
+import torch
+from promise import Promise2dDataset
+import skimage.transform
+import numpy as np
+
+def train_transform(input, target):
+    input = skimage.transform.resize(input, (320, 320), mode='constant')
+    target = skimage.transform.resize(target, (320, 320), mode='constant')
+    if np.random.random() < 0.5:
+        input = np.flip(input, 1).copy()
+        target = np.flip(target, 1).copy()
+    input = torch.from_numpy(input).unsqueeze(0)
+    target = torch.from_numpy(target).unsqueeze(0)
+    return input, target
+
+def valid_transform(input, target):
+    input = skimage.transform.resize(input, (320, 320))
+    target = skimage.transform.resize(target, (320, 320))
+    input = torch.from_numpy(input).unsqueeze(0)
+    target = torch.from_numpy(target).unsqueeze(0)
+    return input, target
+
+trainset = Promise2dDataset('promisedata/train', train=True, threshold=0,
+                            split=0.7, transform=train_transform)
+validset = Promise2dDataset('promisedata/train', train=False, threshold=0,
+                            split=0.7, transform=valid_transform)
+```
+
 # UNet-Zoo
 A collection of UNet and hybrid architectures for 2D and 3D Biomedical Image segmentation, implemented in PyTorch.
 

promise.py

@@ -0,0 +1,40 @@
+import torch.utils.data
+import SimpleITK as sitk
+import numpy as np
+import os
+import skimage.transform
+
+class Promise2dDataset(torch.utils.data.Dataset):
+        def __init__(self, root, threshold = 0, train=True, split= 0.7, transform=None):
+            self.transform = transform
+
+            files = [(os.path.join(root, j), os.path.join(root, j[:-4] + '_segmentation.mhd'))
+                for j in os.listdir(root) if j.endswith('mhd') and len(j) == 10]
+
+            if train:
+                files = files[:int(len(files) * split)]
+            else:
+                files = files[int(len(files) * split):]
+            x, y = [], []
+            for f1, f2 in files:
+                ctscans = sitk.GetArrayFromImage(sitk.ReadImage(f1))
+                segmaps = sitk.GetArrayFromImage(sitk.ReadImage(f2))
+                fractions = np.mean(segmaps, axis=(1,2))
+                for j in range(len(fractions)):
+                    if fractions[j] > threshold:
+                        x.append(ctscans[j])
+                        y.append(segmaps[j])
+
+            self.X = x
+            self.Y = y
+
+        def __len__(self):
+            return len(self.X)
+
+        def __getitem__(self, index):
+            input = self.X[index].astype(float)
+            target = self.Y[index].astype(float)
+            if self.transform:
+                input, target = self.transform(input, target)
+            return input, target
+

test_promise.py

@@ -0,0 +1,46 @@
+from promise import Promise2dDataset
+import torch
+import numpy as np
+import skimage.transform
+import matplotlib.pyplot as plt
+
+def train_transform(input, target):
+    input = skimage.transform.resize(input, (320, 320), mode='constant')
+    target = skimage.transform.resize(target, (320, 320), mode='constant')
+    if np.random.random() < 0.5:
+        input = np.flip(input, 1).copy()
+        target = np.flip(target, 1).copy()
+    input = torch.from_numpy(input).unsqueeze(0)
+    target = torch.from_numpy(target).unsqueeze(0)
+    return input, target
+
+def valid_transform(input, target):
+    input = skimage.transform.resize(input, (320, 320), mode='constant')
+    target = skimage.transform.resize(target, (320, 320), mode='constant')
+    input = torch.from_numpy(input).unsqueeze(0)
+    target = torch.from_numpy(target).unsqueeze(0)
+    return input, target
+
+trainset = Promise2dDataset('promisedata/train', train=True, threshold=0,
+                            split=0.7, transform=train_transform)
+validset = Promise2dDataset('promisedata/train', train=False, threshold=0,
+                            split=0.7, transform=valid_transform)
+
+
+p, q = trainset[0]
+fig, (ax1, ax2) = plt.subplots(1, 2)
+ax1.imshow(p.squeeze(), cmap=plt.cm.gray)
+ax2.imshow(q.squeeze(), cmap=plt.cm.gray)
+plt.show()
+
+train_iter = torch.utils.data.DataLoader(trainset, batch_size=32)
+val_iter = torch.utils.data.DataLoader(validset, batch_size=32)
+
+for p, q in train_iter:
+    print(p.shape, p.max(), p.min())
+    print(q.shape, q.max(), q.min())
+
+
+for p, q in val_iter:
+    print(p.shape, p.max(), p.min())
+    print(q.shape, q.max(), q.min())