Make CLS pipeline

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "yucca"
-version = "2.2.1"
+version = "2.2.4"
 authors = [
   { name="Sebastian Llambias", email="llambias@live.com" },
   { name="Asbjørn Munk", email="9844416+asbjrnmunk@users.noreply.github.com" },
@@ -29,6 +29,7 @@ dependencies = [
     "numpy>=1.26.4",
     "pandas>=2.2.1",
     "python-dotenv==1.0.0",
+    "pytorchvideo==0.1.5",
     "scikit_image>=0.22.0",
     "scikit_learn>=1.4.1.post1",
     "seaborn>=0.13.2",

yucca/functional/evaluation/confusion_matrix.py

@@ -35,3 +35,12 @@ def torch_get_tp_fp_tn_fn(confusion_matrix, ignore_label=0):
         TN.append(tn.cpu().numpy())
         FN.append(fn.cpu().numpy())
     return TP, FP, TN, FN
+
+
+def convert_confusion_matrix_to_dict(confusion_matrix):
+    d = {}
+    for true_label, row in enumerate(confusion_matrix):
+        d[str(true_label)] = {}
+        for predicted_label, value in enumerate(row):
+            d[str(true_label)][str(predicted_label)] = value
+    return d

yucca/functional/evaluation/evaluate_folder.py

@@ -1,4 +1,5 @@
 import sys
+import os
 import numpy as np
 import nibabel as nib
 import logging
@@ -11,6 +12,7 @@
 from batchgenerators.utilities.file_and_folder_operations import join
 from sklearn.metrics import confusion_matrix
 from yucca.functional.evaluation.metrics import auroc
+from yucca.functional.evaluation.confusion_matrix import convert_confusion_matrix_to_dict
 
 
 def evaluate_folder_segm(
@@ -257,9 +259,6 @@ def evaluate_folder_cls(
     prediction_probs = []
     ground_truths = []
 
-    # Flag to check if we have prediction probabilities to calculate AUROC
-    use_probs = False
-
     # load predictions and ground truths
     for case in tqdm(subjects, desc="Evaluating"):
         predpath = join(folder_with_predictions, case)
@@ -268,15 +267,9 @@ def evaluate_folder_cls(
         pred: int = np.loadtxt(predpath)
         gt: int = np.loadtxt(gtpath)
 
-        try:
-            if len(prediction_probs) == 0:
-                print("Prediction probabilities found. Will use them for evaluation.")
-                use_probs = True
-
+        if os.path.isfile(predpath.replace(".txt", ".npz")):
             pred_probs = np.load(predpath.replace(".txt", ".npz"))["data"]  # contains output probabilities
             prediction_probs.append(pred_probs)
-        except FileNotFoundError:
-            pred_probs = None
 
         predictions.append(pred)
         ground_truths.append(gt)
@@ -287,6 +280,8 @@ def evaluate_folder_cls(
     # calculate per-class metrics
     cmat = confusion_matrix(ground_truths, predictions, labels=labels)
 
+    cmat_dict = convert_confusion_matrix_to_dict(cmat)
+    resultdict["confusion_matrix"] = cmat_dict
     resultdict["per_class"] = {}
 
     for label in labels:
@@ -303,7 +298,7 @@ def evaluate_folder_cls(
         resultdict["per_class"][str(label)] = labeldict
 
     # calculate AUROC
-    if use_probs:
+    if len(prediction_probs) > 0:
         auroc_per_class: list[float] = auroc(ground_truths, prediction_probs)
         for label, score in zip(labels, auroc_per_class):
             resultdict["per_class"][str(label)]["AUROC"] = round(score, 4)

yucca/functional/preprocessing.py

@@ -402,6 +402,7 @@ def preprocess_case_for_inference(
     target_size,
     target_spacing,
     target_orientation,
+    background_pixel_value: int = 0,
     allow_missing_modalities: bool = False,
     ext=".nii.gz",
     keep_aspect_ratio: bool = True,
@@ -425,7 +426,9 @@ def preprocess_case_for_inference(
     image_properties["uncropped_shape"] = np.array(images[0].shape)
 
     if crop_to_nonzero:
-        nonzero_box = get_bbox_for_foreground(images[0], background_label=0)
+        if np.max(images[0]) <= background_pixel_value:
+            background_pixel_value = np.min(images[0])
+        nonzero_box = get_bbox_for_foreground(images[0], background_label=background_pixel_value)
         for i in range(len(images)):
             images[i] = crop_to_box(images[i], nonzero_box)
         image_properties["nonzero_box"] = nonzero_box

yucca/functional/testing/data/nifti.py

@@ -1,6 +1,7 @@
 import nibabel as nib
 import numpy as np
 import nibabel.orientations as nio
+import logging
 from yucca.functional.utils.nib_utils import get_nib_orientation, get_nib_spacing
 
 
@@ -48,4 +49,5 @@ def verify_orientation_is_LR_PA_IS(image: nib.Nifti1Image):
     if np.all(nio.axcodes2ornt(orientation)[:, 0] == expected_orientation_code):
         return True
     else:
+        logging.info(f"Found orientation {orientation}")
         return False

yucca/functional/transforms/gamma.py

@@ -9,6 +9,7 @@ def augment_gamma(
     invert_image=False,
     epsilon=1e-7,
     per_channel=False,
+    p_per_channel=None,
     clip_to_input_range=False,
 ):
     if invert_image:
@@ -27,19 +28,20 @@ def augment_gamma(
             data_sample = np.clip(data_sample, a_min=img_min, a_max=img_max)
     else:
         for c in range(data_sample.shape[0]):
-            if np.random.random() < 0.5 and gamma_range[0] < 1:
-                gamma = np.random.uniform(gamma_range[0], 1)
-            else:
-                gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
-            img_min = data_sample[c].min()
-            img_max = data_sample[c].max()
-            img_range = img_max - img_min
-            data_sample[c] = (
-                np.power(((data_sample[c] - img_min) / float(img_range + epsilon)), gamma) * float(img_range + epsilon)
-                + img_min
-            )
-            if clip_to_input_range:
-                data_sample[c] = np.clip(data_sample[c], a_min=img_min, a_max=img_max)
+            if np.random.uniform() < p_per_channel[c]:
+                if np.random.random() < 0.5 and gamma_range[0] < 1:
+                    gamma = np.random.uniform(gamma_range[0], 1)
+                else:
+                    gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
+                img_min = data_sample[c].min()
+                img_max = data_sample[c].max()
+                img_range = img_max - img_min
+                data_sample[c] = (
+                    np.power(((data_sample[c] - img_min) / float(img_range + epsilon)), gamma) * float(img_range + epsilon)
+                    + img_min
+                )
+                if clip_to_input_range:
+                    data_sample[c] = np.clip(data_sample[c], a_min=img_min, a_max=img_max)
     if invert_image:
         data_sample = -data_sample
     return data_sample

yucca/functional/transforms/spatial.py

@@ -26,6 +26,7 @@ def spatial(
     scale_factor,
     clip_to_input_range,
     label: Optional[np.ndarray] = None,
+    linear_interpolation_channel=None,
     skip_label: bool = False,
     do_crop: bool = True,
     random_crop: bool = True,
@@ -40,6 +41,9 @@ def spatial(
         cval = cval
     assert isinstance(cval, (int, float)), f"got {cval} of type {type(cval)}"
 
+    if isinstance(order, (int, float)):
+        order = [order for _ in range(image.shape[1])]
+
     coords = create_zero_centered_coordinate_matrix(patch_size)
     image_canvas = np.zeros((image.shape[0], image.shape[1], *patch_size), dtype=np.float32)
 
@@ -92,7 +96,7 @@ def spatial(
             image_canvas[b, c] = map_coordinates(
                 image[b, c].astype(float),
                 coords,
-                order=order,
+                order=order[c],
                 mode="constant",
                 cval=cval,
             ).astype(image.dtype)
@@ -106,7 +110,7 @@ def spatial(
             dtype=np.float32,
         )
 
-        # Mapping the labelmentations to the distorted coordinates
+        # Mapping the label to the distorted coordinates
         for b in range(label.shape[0]):
             for c in range(label.shape[1]):
                 label_canvas[b, c] = map_coordinates(label[b, c], coords, order=0, mode="constant", cval=0.0).astype(

yucca/functional/visualization/__init__.py

@@ -1 +1,3 @@
-from yucca.functional.visualization.imshow import get_train_fig_with_inp_out_tar
+from yucca.functional.visualization.imshow import get_segm_train_fig_with_inp_out_tar
+from yucca.functional.visualization.imshow import get_cls_train_fig_with_inp_out_tar
+from yucca.functional.visualization.imshow import get_ssl_train_fig_with_inp_out_tar

yucca/functional/visualization/imshow.py

@@ -3,18 +3,14 @@
 import matplotlib.pyplot as plt
 
 
-def get_train_fig_with_inp_out_tar(input, output, target, fig_title, task_type: str = "segmentation"):
+def get_segm_train_fig_with_inp_out_tar(input, output, target, fig_title):
     # This needs to handle the following cases:
     # Segmentation      : {"input": (m,x,y(,z)), "target": (1,x,y(,z)), "output": (c,x,y(,z))}
-    # Self-supervised   : {"input": (m,x,y(,z)), "target": (m,x,y(,z)), "output": (m,x,y(,z))}
-    # Classification    : {"input": (m,x,y(,z)), "target": (1,x), "output": (c,x)}
 
     channel_idx = np.random.randint(0, input.shape[0])
 
-    if len(input.shape) == 4:  # 3D images.
-        # We need to select a slice to visualize.
-        if task_type == "segmentation" and len(target[0].nonzero()[0]) > 0:
-            # Select a foreground slice if any exist.
+    if len(input.shape) == 4:
+        if len(target[0].nonzero()[0]) > 0:
             foreground_locations = target[0].nonzero()
             slice_to_visualize = foreground_locations[0][np.random.randint(0, len(foreground_locations[0]))]
         else:
@@ -27,30 +23,67 @@ def get_train_fig_with_inp_out_tar(input, output, target, fig_title, task_type:
             output = output[:, slice_to_visualize]
 
     image = input[channel_idx]
+    target = target[0]
+    output = output.argmax(0)
+
+    fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5), dpi=100, constrained_layout=True)
+    axes[0].imshow(image, cmap="gray", vmin=np.quantile(image, 0.01), vmax=np.quantile(image, 0.99))
+    axes[0].set_title("input")
+    axes[1].imshow(target, cmap="gray")
+    axes[1].set_title("target")
+    axes[2].imshow(output, cmap="gray")
+    axes[2].set_title("output")
+    fig.suptitle(fig_title, fontsize=16)
+    return fig
+
+
+def get_cls_train_fig_with_inp_out_tar(input, output, target, fig_title):
+    # This needs to handle the following case:
+    # Classification    : {"input": (m,x,y(,z)), "target": (n_classes), "output": (n_classes)}
+
+    channel_idx = np.random.randint(0, input.shape[0])
+
+    slice_to_visualize = np.random.randint(0, input.shape[1])
+
+    if len(input.shape) == 4:  # 3D images.
+        input = input[:, slice_to_visualize]
+
+    image = input[channel_idx]
+
+    output = output.argmax(0)
+
+    fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(5, 5), dpi=100, constrained_layout=True)
+    axes.imshow(image, cmap="gray", vmin=np.quantile(image, 0.01), vmax=np.quantile(image, 0.99))
+    axes.set_title(f"Input: {fig_title}", fontsize=12)
+    fig.suptitle(f"Target: {target} | Output: {output}", fontsize=12)
+    return fig
+
+
+def get_ssl_train_fig_with_inp_out_tar(input, output, target, fig_title):
+    # This needs to handle the following cases:
+    # Self-supervised   : {"input": (m,x,y(,z)), "target": (m,x,y(,z)), "output": (m,x,y(,z))}
+
+    channel_idx = np.random.randint(0, input.shape[0])
+
+    if len(input.shape) == 4:  # 3D images.
+        slice_to_visualize = np.random.randint(0, input.shape[1])
+        input = input[:, slice_to_visualize]
+        if len(target.shape) == 4:
+            target = target[:, slice_to_visualize]
+        if len(output.shape) == 4:
+            output = output[:, slice_to_visualize]
+
+    image = input[channel_idx]
+
+    target = target[channel_idx]
+    output = output[channel_idx]
 
-    if task_type in ["segmentation", "classification"]:
-        target = target[0]
-        output = output.argmax(0)
-    elif task_type == "self-supervised":
-        target = target[channel_idx]
-        output = output[channel_idx]
-    else:
-        logging.warn(
-            f"Unknown task type. Found {task_type} and expected one in ['classification', 'segmentation', 'self-supervised']"
-        )
-
-    if len(target.shape) == 1:
-        fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(5, 5), dpi=100, constrained_layout=True)
-        axes[0].imshow(image, cmap="gray", vmin=np.quantile(image, 0.01), vmax=np.quantile(image, 0.99))
-        axes[0].set_title("input")
-        fig.suptitle(f"{fig_title}. Target: {target} | Output: {output}", fontsize=16)
-    else:
-        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5), dpi=100, constrained_layout=True)
-        axes[0].imshow(image, cmap="gray", vmin=np.quantile(image, 0.01), vmax=np.quantile(image, 0.99))
-        axes[0].set_title("input")
-        axes[1].imshow(target, cmap="gray")
-        axes[1].set_title("target")
-        axes[2].imshow(output, cmap="gray")
-        axes[2].set_title("output")
-        fig.suptitle(fig_title, fontsize=16)
+    fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5), dpi=100, constrained_layout=True)
+    axes[0].imshow(image, cmap="gray", vmin=np.quantile(image, 0.01), vmax=np.quantile(image, 0.99))
+    axes[0].set_title("input")
+    axes[1].imshow(target, cmap="gray")
+    axes[1].set_title("target")
+    axes[2].imshow(output, cmap="gray")
+    axes[2].set_title("output")
+    fig.suptitle(fig_title, fontsize=16)
     return fig

yucca/modules/callbacks/loggers.py

@@ -74,7 +74,7 @@ def create_logfile(self):
             self.log_dir,
             "training_log.txt",
         )
-        with open(self.log_file, "w") as f:
+        with open(self.log_file, "a+") as f:
             f.write("Starting model training")
             logging.info("Starting model training \n" f'{"log file:":20} {self.log_file} \n')
             f.write("\n")