 b/adpkd_segmentation/datasets/masks.py
+import numpy as np
+BACKGROUND = 0.0
+L_KIDNEY = 0.5019608
+R_KIDNEY = 0.7490196
+BACKGROUND_INT = 0
+L_KIDNEY_INT = 128
+R_KIDNEY_INT = 191
+class SingleChannelMaskNumpy:
+    """Sets 1 for kidneys, 0 otherwise."""
+    def __call__(self, label):
+        """
+        Args:
+            label, (1, H, W) uint8 numpy array
+        Returns:
+            numpy array, (1, H, W) uint8 one-hot encoded mask
+        """
+        kidney = np.bitwise_or(label == R_KIDNEY_INT, label == L_KIDNEY_INT)
+        return kidney.astype(np.uint8)
+class TwoChannelsMaskNumpy:
+    """
+    The first channel for right kidney vs background,
+    and the second one for left kidney vs background.
+    Kidneys are marked as 1, background as 0.
+    """
+    def __call__(self, label):
+        """
+        Args:
+            label, (1, H, W) uint8 numpy array
+        Returns:
+            numpy array, (2, H, W) uint8 one-hot encoded mask
+        """
+        r_kidney = label == R_KIDNEY_INT
+        l_kidney = label == L_KIDNEY_INT
+        mask = np.concatenate([r_kidney, l_kidney], axis=0).astype(np.uint8)
+        return mask
+class ThreeChannelMaskNumpy:
+    """One channel for each of the 3 classes. Background last."""
+    def __call__(self, label):
+        """
+        Args:
+            label, (1, H, W) float32 tensor
+        Returns:
+            numpy array, (3, H, W) uint8 one-hot encoded mask
+        """
+        background = label == BACKGROUND_INT
+        r_kidney = label == R_KIDNEY_INT
+        l_kidney = label == L_KIDNEY_INT
+        mask = np.concatenate([r_kidney, l_kidney, background], axis=0).astype(
+            np.uint8
+        )
+        return mask