make exercise exercisable

Author: loki-veera

README.md

@@ -7,18 +7,21 @@ from T2-Weighted MRI"](https://www.var.ovgu.de/pub/2019_Meyer_ISBI_Zone_Segmenta
 by Meyer et al.
 
 
-1. To get started run
+### Task 1: To get started run
 
 ```bash
 python ./data/download.py
 ```
 
 in your terminal. The script will download and prepare the medical scans and domain-expert
-annotations for you.
+annotations for you or you can copy the data from bender at the following location.
+```bash
+TODO: update bender location here.
+```
 
-Data loading and resampling work already. 
+Data loading and resampling work already. The next task is optional. If you want to skip it, download the `compute_roi.py` from eCampus and replace the contents with the existing function `compute_roi()` in the repository.
 
-1. #### Find the bounding box roi as described below by finishing the `compute_roi` function.
+### Task 2 (Optional): Find the bounding box roi as described below by finishing the `compute_roi` function. 
 Once you have obtained the train and test data, you must create a preprocessing pipeline.
 Proceed to `src/util.py` and compute the so called region of interest.
 Meyer et al. define this region as:
@@ -70,32 +73,27 @@ local coordinates now allows array indexing. Following Meyer et al. we discard a
 
 Test your implementation by setting the if-condition wrapping the plotting utility in `compute_roi` to `True` and running vscode pytest `test_roi`. Remember to set it back to `False` afterwards.
 
-2. #### Implement the UNet. 
+### Task 3: Implement the UNet. 
+Navigate to the `train.py` file in the `src` folder.
+Finish the `UNet3D` class, as discussed in the lecture.
+Use [torch.nn.Conv3d](https://pytorch.org/docs/stable/generated/torch.nn.Conv3d.html), [torch.nn.ReLU](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html), [torch.nn.MaxPool3d](https://pytorch.org/docs/stable/generated/torch.nn.MaxPool3d.html) and [th.nn.UpSample](https://pytorch.org/docs/stable/generated/torch.nn.Upsample.html) to build the model. For upsampling, we suggest to use `mode='nearest'` algorithm for reproducibility purpose.
 
-Navigate to the `train.py` module file in the `src` folder. 
-Finish the `UNet3D` class, as discussed in the lecture. 
-Use the [flax.linen.Conv](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Conv.html), [flax.linen.relu](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.activation.relu.html), and [flax.linen.ConvTranspose](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.ConvTranspose.html), to build your model.
-
-3. #### Implement the focal-loss
+### Task 4: Implement the focal-loss.
 
 Open the `util.py` module in `src` and implement the `softmax_focal_loss` function as discussed in the lecture:
 
 $$\mathcal{L}(\mathbf{o},\mathbf{I})=-\mathbf{I}\cdot(1-\sigma_s(\mathbf{o}))^\gamma\cdot\alpha\cdot\ln(\sigma_s(\mathbf{o})) $$
 
 with output logits $\mathbf{o}$, the corresponding labels $\mathbf{I}$ and the softmax function $\sigma_s$.
 
-4. #### Run and test the training script.
+### Task 5: Run and test the training script.
 
 Execute the training script with by running `scripts/train.slurm` (locally or using `sbatch`).
 
 After training you can test your model by changing the `checkpoint_name` variable in `src/sample.py` to the desired model checkpoint and running `scripts/test.slurm`.
 
-#### Solution:
-![slice](./fig/prostatext2.png)
-![slice](./fig/prostatext2_net.png)
-![slice](./fig/prostatext2_true.png)
 
-5. #### (Optional) Implement mean Intersection-over-Union (mIoU)
+### Task 6: Implement mean Intersection-over-Union (mIoU)
 
 Open the `meanIoU.py` in `src` and implement the `compute_iou` function as discussed below.
 mIoU is the most common metric used for evaluating semantic segmentation tasks. It can be computed using the values from a confusion matrix as given below
@@ -113,5 +111,4 @@ python -m src.meanIoU
 
 ### Acknowledgments:
 We thank our course alumni Barbara Wichtmann, for bringing this problem to our attention.
-Without her feedback, this code would not exist.
-
+Without her feedback, this code would not exist.

README.pdf

@@ -19,26 +19,8 @@ def compute_iou(preds: th.Tensor, target: th.Tensor) -> th.Tensor:
         jnp.ndarray: Mean Intersection over Union values
     """
     assert preds.shape == target.shape
-
-    b, h, w, s = target.shape
-    preds = preds.permute((0, 3, 1, 2))
-    target = target.permute((0, 3, 1, 2))
-    batch_preds = th.reshape(preds, (b * s, h, w, 1))
-    batch_target = th.reshape(target, (b * s, h, w, 1))
-    batch_iou = []
-    for idx in range(b * s):
-        preds = batch_preds[idx]
-        target = batch_target[idx]
-        per_class_iou = []
-        for cls in range(0, 5):
-            if th.any(preds == cls) or th.any(target == cls):
-                tp = th.sum((preds == cls) & (target == cls))
-                fp = th.sum((preds != cls) & (target == cls))
-                fn = th.sum((preds == cls) & (target != cls))
-                iou = tp / (tp + fp + fn + 1e-8)
-                per_class_iou.append(iou)
-        batch_iou.append(th.mean(th.tensor(per_class_iou)))
-    return th.mean(th.tensor(batch_iou))
+    # TODO: Implement meanIoU
+    return th.tensor(0.0)
 
 
 if __name__ == "__main__":

src/meanIoU.py

@@ -74,79 +74,8 @@ def __init__(self):
         input_feat = 1
         init_feat = 16
         out_neurons = 5
-        # Five Downscale blocks
-        self.downscale_1 = th.nn.Sequential(
-            th.nn.Conv3d(input_feat, init_feat, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat, init_feat, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.downscale_2 = th.nn.Sequential(
-            th.nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2)),
-            # th.nn.BatchNorm3d(init_feat),
-            th.nn.Conv3d(init_feat, init_feat * 2, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 2, init_feat * 2, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.downscale_3 = th.nn.Sequential(
-            th.nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2)),
-            # th.nn.BatchNorm3d(init_feat * 2),
-            th.nn.Conv3d(init_feat * 2, init_feat * 4, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 4, init_feat * 4, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.downscale_4 = th.nn.Sequential(
-            th.nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2)),
-            # th.nn.BatchNorm3d(init_feat * 4),
-            th.nn.Conv3d(init_feat * 4, init_feat * 8, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 8, init_feat * 8, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.downscale_5 = th.nn.Sequential(
-            th.nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2)),
-            # th.nn.BatchNorm3d(init_feat * 8),
-            th.nn.Conv3d(init_feat * 8, init_feat * 16, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 16, init_feat * 16, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        # Four Upscale conv blocks
-        self.upscale_4 = th.nn.Sequential(
-            # th.nn.BatchNorm3d(init_feat * 16  + init_feat * 8),
-            th.nn.Conv3d(
-                init_feat * 16 + init_feat * 8, init_feat * 8, (3, 3, 3), padding=1
-            ),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 8, init_feat * 8, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.upscale_3 = th.nn.Sequential(
-            # th.nn.BatchNorm3d(init_feat * 8 + init_feat * 4),
-            th.nn.Conv3d(
-                init_feat * 8 + init_feat * 4, init_feat * 4, (3, 3, 3), padding=1
-            ),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 4, init_feat * 4, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.upscale_2 = th.nn.Sequential(
-            # th.nn.BatchNorm3d(init_feat * 4 + init_feat * 2),
-            th.nn.Conv3d(
-                init_feat * 4 + init_feat * 2, init_feat * 2, (3, 3, 3), padding=1
-            ),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat * 2, init_feat * 2, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-        )
-        self.upscale_1 = th.nn.Sequential(
-            # th.nn.BatchNorm3d(init_feat * 2 + init_feat),
-            th.nn.Conv3d(init_feat * 2 + init_feat, init_feat, (3, 3, 3), padding=1),
-            th.nn.ReLU(),
-            th.nn.Conv3d(init_feat, out_neurons, (3, 3, 3), padding=1),
-        )
+        # TODO: Initialize downscaling blocks
+        # TODO: Initialize upscaling blocks
 
     def forward(self, x: th.Tensor) -> th.Tensor:
         """Forward pass.
@@ -157,42 +86,8 @@ def forward(self, x: th.Tensor) -> th.Tensor:
         Returns:
             th.Tensor: Segmented output.
         """
-        x1 = self.downscale_1(x)
-        x1 = pad_odd(x1)
-
-        x2 = self.downscale_2(x1)
-        x2 = pad_odd(x2)
-
-        x3 = self.downscale_3(x2)
-        x3 = pad_odd(x3)
-
-        x4 = self.downscale_4(x3)
-        x4 = pad_odd(x4)
-
-        x5 = self.downscale_5(x4)
-        x5 = pad_odd(x5)
-
-        x6 = self.__upsize(x5)
-        x6 = x6[:, :, : x4.shape[2], : x4.shape[3], : x4.shape[4]]
-        x6 = th.cat([x4, x6], dim=1)
-        x6 = self.upscale_4(x6)
-
-        x7 = self.__upsize(x6)
-        x7 = x7[:, :, : x3.shape[2], : x3.shape[3], : x3.shape[4]]
-        x7 = th.cat([x3, x7], dim=1)
-        x7 = self.upscale_3(x7)
-
-        x8 = self.__upsize(x7)
-        x8 = x8[:, :, : x2.shape[2], : x2.shape[3], : x2.shape[4]]
-        x8 = th.cat([x2, x8], dim=1)
-        x8 = self.upscale_2(x8)
-
-        x9 = self.__upsize(x8)
-        x9 = x9[:, :, : x1.shape[2], : x1.shape[3], : x1.shape[4]]
-        x9 = th.cat([x1, x9], dim=1)
-        x9 = self.upscale_1(x9)
-        out = x9[:, :, : x.shape[2], : x.shape[3], : x.shape[4]]
-        return out
+        # TODO: Implement 3D UNet as discussed in the lecture
+        return th.tensor(0.0)
 
     def __upsize(self, input_: th.Tensor) -> th.Tensor:
         """Upsample image.
@@ -203,8 +98,8 @@ def __upsize(self, input_: th.Tensor) -> th.Tensor:
         Returns:
             th.Tensor: Upsampled image.
         """
-        _, _, d, h, w = input_.shape
-        return th.nn.Upsample(size=(d, h * 2, w * 2), mode="nearest")(input_)
+        # TODO: Upsample the height and width using th.nn.Upsample with nearest mode.
+        return th.tensor(0.0)
 
 
 def train():
@@ -225,7 +120,7 @@ def train():
 
     model = UNet3D().to(device)
     opt = th.optim.Adam(model.parameters(), lr=1e-4)
-    load_new = False
+    load_new = True
 
     writer = metric_writers.create_default_writer(
         "./runs/" + str(datetime.now()), asynchronous=False
@@ -251,7 +146,6 @@ def train():
     val_loss_list = []
     train_loss_lost = []
     iter_count = 0
-    # loss_fn = th.nn.CrossEntropyLoss()
 
     for e in range(epochs):
         random.shuffle(epoch_batches)

src/train.py

@@ -2,14 +2,11 @@
 
 from typing import List, Tuple
 
-import chex
-import jax
-import jax.numpy as jnp
 import matplotlib.colors as mcolors
 import matplotlib.pyplot as plt
 import numpy as np
-import optax
 import SimpleITK as sitk  # noqa: N813
+import torch as th
 from SimpleITK.SimpleITK import Image
 
 # from . import zone_segmentation_utils as utils
@@ -140,20 +137,26 @@ def compute_roi(images: Tuple[Image, Image, Image]):
     rects = []
     for pos, size in enumerate(sizes):
         lines = box_lines(size)
-        rotated = [(rotation[pos] @ line.T).T for line in lines]
-        shifted = [origins[pos] + line for line in rotated]
+        # TODO: Rotate and shift the lines.
+        rotated = []
+        shifted = []
         rects.append(shifted)
 
     # find the intersection.
     rects_stacked = np.stack(rects)  # Had to rename because of mypy
+    # TODO: Find the axis maxima and minima
     bbs = [
-        (np.amin(rect, axis=(0, 1)), np.amax(rect, axis=(0, 1)))
+        (
+            np.zeros_like(rect[0, 0]),
+            np.zeros_like(rect[0, 0]),
+        )  # TODO: fixme
         for rect in rects_stacked
     ]
 
     # compute intersection
-    lower_end = np.amax(np.stack([bb[0] for bb in bbs], axis=0), axis=0)
-    upper_end = np.amin(np.stack([bb[1] for bb in bbs], axis=0), axis=0)
+    # TODO: Implement me.
+    lower_end = np.zeros_like(bbs[0][0])
+    upper_end = np.zeros_like(bbs[0][1])
     roi_bb = np.stack((lower_end, upper_end))
     roi_bb_size = roi_bb[1] - roi_bb[0]
 
@@ -164,8 +167,8 @@ def compute_roi(images: Tuple[Image, Image, Image]):
     # compute roi coordinates in image space.
     img_coord_rois = [
         (
-            (np.linalg.inv(rot) @ (roi_bb[0] - offset).T).T / spacing,
-            (np.linalg.inv(rot) @ (roi_bb[1] - offset).T).T / spacing,
+            np.zeros_like(roi_bb[0]),  # TODO: Implement me
+            np.zeros_like(roi_bb[1]),  # TODO: Implement me
         )
         for rot, offset, spacing in zip(rotation, origins, spacings)
     ]
@@ -244,54 +247,6 @@ def in_array(in_int, dim):
     return intersections, box_indices
 
 
-def sigmoid_focal_loss(
-    logits: jnp.ndarray,
-    labels: jnp.ndarray,
-    alpha: float = -1,
-    gamma: float = 2,
-) -> jnp.ndarray:
-    """Compute a sigmoid focal loss.
-
-    Implementation of the focal loss as used https://arxiv.org/abs/1708.02002.
-    This loss often appears in the segmentation context.
-    Use this loss function if classes are not mutually exclusive.
-    See `sigmoid_binary_cross_entropy` for more information.
-
-    Args:
-        logits: A float array of arbitrary shape.
-                The predictions for each example.
-        labels: A float array, its shape must be identical to
-                that of logits. It containes the binary
-                 classification label for each element in logits
-                (0 for the out of class and 1 for in class).
-                This array is often one-hot encoded.
-        alpha: (optional) Weighting factor in range (0,1) to balance
-                positive vs negative examples. Default = -1 (no weighting).
-        gamma: Exponent of the modulating factor (1 - p_t) to
-               balance easy vs hard examples.
-
-    Returns:
-        A loss value array with a shape identical to the logits and target
-        arrays.
-    """
-    chex.assert_type([logits], float)
-    labels = labels.astype(logits.dtype)
-
-    # see also the original implementation at:
-    # https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/focal_loss.py
-    p = jax.nn.sigmoid(logits)
-    ce_loss = optax.sigmoid_binary_cross_entropy(logits, labels)
-    p_t = p * labels + (1 - p) * (1 - labels)
-    loss = ce_loss * ((1 - p_t) ** gamma)
-    if alpha >= 0:
-        alpha_t = alpha * labels + (1 - alpha) * (1 - labels)
-        loss = alpha_t * loss
-    return loss
-
-
-import torch as th
-
-
 def softmax_focal_loss(
     logits: th.Tensor,
     labels: th.Tensor,
@@ -308,25 +263,5 @@ def softmax_focal_loss(
     # return jnp.sum(loss, axis=-1)
     logits = logits.float()
     labels = labels.float()
-    focus = th.pow(1.0 - th.nn.functional.softmax(logits, dim=-1), gamma)
-    loss = -labels * focus * alpha * th.nn.functional.log_softmax(logits, dim=-1)
-    return th.sum(loss, dim=-1)
-
-
-# def tversky(y_true, y_pred, alpha=.3, beta=.7):
-#     """See: https://arxiv.org/pdf/1706.05721.pdf"""
-#     y_true_f = jnp.reshape(y_true, -1)
-#     y_pred_f = jnp.reshape(y_pred, -1)
-#     intersection = jnp.sum(y_true_f * y_pred_f)
-#     G_P = alpha * jnp.sum((1 - y_true_f) * y_pred_f)  # G not P
-#     P_G = beta * jnp.sum(y_true_f * (1 - y_pred_f))  # P not G
-#     return (intersection + 1.) / (intersection + 1. + G_P + P_G)
-#
-# def Tversky_loss(y_true, y_pred):
-#     return -tversky(y_true, y_pred)
-#
-#
-# def dice_coeff(logits, labels):
-#     pred_probs = jax.nn.softmax(logits)
-#     intersection = jnp.sum(labels * logits)
-#     return ((2. * intersection + 1.) / (jnp.sum(labels) + jnp.sum(pred_probs) + 1.))*(-1.)
+    # TODO: Implement softmax focal loss.
+    return th.tensor(0.0)