* fixes consecutive slices 2d & 3D settings for mtlrs

Author: wdika

atommic/collections/multitask/rs/data/mrirs_loader.py

@@ -409,26 +409,32 @@ def __getitem__(self, i: int):  # noqa: MC0001
             kspace = self.get_consecutive_slices(hf, "kspace", dataslice).astype(np.complex64)
 
             if not is_none(dataset_format) and dataset_format == "skm-tea-echo1":
-                kspace = kspace[:, :, 0, :]
+                kspace = kspace[..., 0, :]
             elif not is_none(dataset_format) and dataset_format == "skm-tea-echo2":
-                kspace = kspace[:, :, 1, :]
+                kspace = kspace[..., 1, :]
             elif not is_none(dataset_format) and dataset_format == "skm-tea-echo1+echo2":
-                kspace = kspace[:, :, 0, :] + kspace[:, :, 1, :]
+                kspace = kspace[..., 0, :] + kspace[..., 1, :]
             elif not is_none(dataset_format) and dataset_format == "skm-tea-echo1+echo2-mc":
-                kspace = np.concatenate([kspace[:, :, 0, :], kspace[:, :, 1, :]], axis=-1)
+                kspace = np.concatenate([kspace[..., 0, :], kspace[..., 1, :]], axis=-1)
             else:
                 warnings.warn(
                     f"Dataset format {dataset_format} is either not supported or set to None. "
                     "Using by default only the first echo."
                 )
-                kspace = kspace[:, :, 0, :]
-
-            kspace = kspace[48:-48, 40:-40]
+                kspace = kspace[..., 0, :]
 
             sensitivity_map = self.get_consecutive_slices(hf, "maps", dataslice).astype(np.complex64)
             sensitivity_map = sensitivity_map[..., 0]
 
-            sensitivity_map = sensitivity_map[48:-48, 40:-40]
+            if dataset_format == "skm-tea-echo1+echo2-mc":
+                sensitivity_map = np.concatenate([sensitivity_map, sensitivity_map], axis=-1)
+
+            if self.consecutive_slices > 1:
+                sensitivity_map = sensitivity_map[:, 48:-48, 40:-40]
+                kspace = kspace[:, 48:-48, 40:-40]
+            else:
+                sensitivity_map = sensitivity_map[48:-48, 40:-40]
+                kspace = kspace[48:-48, 40:-40]
 
             if masking == "custom":
                 mask = np.array([])
@@ -484,12 +490,11 @@ def __getitem__(self, i: int):  # noqa: MC0001
             # TODO: This is hardcoded on the SKM-TEA side, how to generalize this?
             # We need to crop the segmentation labels in the frequency domain to reduce the FOV.
             segmentation_labels = np.fft.fftshift(np.fft.fft2(segmentation_labels))
-            segmentation_labels = segmentation_labels[:, 48:-48, 40:-40]
+            segmentation_labels = segmentation_labels[..., 48:-48, 40:-40]
             segmentation_labels = np.fft.ifft2(np.fft.ifftshift(segmentation_labels)).real
             segmentation_labels = np.where(segmentation_labels > 0.5, 1.0, 0.0)  # Make sure the labels are binary.
 
             imspace = np.empty([])
-
             initial_prediction = np.empty([])
             attrs = dict(hf.attrs)
 
@@ -501,8 +506,12 @@ def __getitem__(self, i: int):  # noqa: MC0001
 
             attrs.update(metadata)
 
-        kspace = np.transpose(kspace, (2, 0, 1))
-        sensitivity_map = np.transpose(sensitivity_map.squeeze(), (2, 0, 1))
+        if self.consecutive_slices == 1:
+            kspace = np.transpose(kspace, (2, 0, 1))
+            sensitivity_map = np.transpose(sensitivity_map.squeeze(), (2, 0, 1))
+        else:
+            kspace = np.transpose(kspace, (0, 3, 1, 2))
+            sensitivity_map = np.transpose(sensitivity_map.squeeze(), (0, 3, 1, 2))
 
         attrs["log_image"] = bool(dataslice in self.indices_to_log)
 

atommic/collections/multitask/rs/nn/base.py

@@ -726,6 +726,9 @@ def __compute_loss__(
         if self.consecutive_slices > 1:
             batch_size, slices = target_segmentation.shape[:2]
             target_segmentation = target_segmentation.reshape(batch_size * slices, *target_segmentation.shape[2:])
+            predictions_segmentation = predictions_segmentation.reshape(  # type: ignore
+                batch_size * slices, *predictions_segmentation.shape[2:]  # type: ignore
+            )
 
         segmentation_loss = self.process_segmentation_loss(target_segmentation, predictions_segmentation, attrs)
 
@@ -798,26 +801,6 @@ def __compute_and_log_metrics_and_outputs__(  # noqa: MC0001
             while isinstance(predictions_segmentation, list):
                 predictions_segmentation = predictions_segmentation[-1]
 
-        if self.consecutive_slices > 1:
-            # reshape the target and prediction to [batch_size, self.consecutive_slices, nr_classes, n_x, n_y]
-            batch_size = target_segmentation.shape[0] // self.consecutive_slices
-            target_segmentation = target_segmentation.reshape(
-                batch_size, self.consecutive_slices, *target_segmentation.shape[1:]
-            )
-            target_reconstruction = target_reconstruction.reshape(
-                batch_size, self.consecutive_slices, *target_reconstruction.shape[2:]
-            )
-            predictions_segmentation = predictions_segmentation.reshape(
-                batch_size, self.consecutive_slices, *predictions_segmentation.shape[2:]
-            )
-            predictions_reconstruction = predictions_reconstruction.reshape(
-                batch_size, self.consecutive_slices, *predictions_reconstruction.shape[1:]
-            )
-            target_segmentation = target_segmentation[:, self.consecutive_slices // 2]
-            target_reconstruction = target_reconstruction[:, self.consecutive_slices // 2]
-            predictions_segmentation = predictions_segmentation[:, self.consecutive_slices // 2]
-            predictions_reconstruction = predictions_reconstruction[:, self.consecutive_slices // 2]
-
         fname = attrs["fname"]
         slice_idx = attrs["slice_idx"]
 
@@ -1267,33 +1250,74 @@ def inference_step(  # noqa: MC0001
             for class_idx, thres in enumerate(self.segmentation_classes_thresholds):
                 if self.segmentation_activation == "sigmoid":
                     if isinstance(predictions_segmentation, list):
-                        cond = [torch.sigmoid(pred[:, class_idx]) for pred in predictions_segmentation]
+                        cond = [
+                            torch.sigmoid(pred[:, class_idx])
+                            if (self.consecutive_slices == 1 or self.dimensionality == 2)
+                            else torch.sigmoid(pred[:, :, class_idx])
+                            for pred in predictions_segmentation
+                        ]
                     else:
-                        cond = torch.sigmoid(predictions_segmentation[:, class_idx])
+                        cond = (
+                            torch.sigmoid(predictions_segmentation[:, class_idx])
+                            if (self.consecutive_slices == 1 or self.dimensionality == 2)
+                            else torch.sigmoid(predictions_segmentation[:, :, class_idx])
+                        )
                 elif self.segmentation_activation == "softmax":
                     if isinstance(predictions_segmentation, list):
-                        cond = [torch.softmax(pred[:, class_idx], dim=1) for pred in predictions_segmentation]
+                        cond = [
+                            torch.softmax(pred[:, class_idx], dim=1)
+                            if (self.consecutive_slices == 1 or self.dimensionality == 2)
+                            else torch.softmax(pred[:, :, class_idx], dim=1)
+                            for pred in predictions_segmentation
+                        ]
                     else:
-                        cond = torch.softmax(predictions_segmentation[:, class_idx], dim=1)
+                        cond = (
+                            torch.softmax(predictions_segmentation[:, class_idx], dim=1)
+                            if (self.consecutive_slices == 1 or self.dimensionality == 2)
+                            else torch.softmax(predictions_segmentation[:, :, class_idx], dim=1)
+                        )
                 else:
                     if isinstance(predictions_segmentation, list):
-                        cond = [pred[:, class_idx] for pred in predictions_segmentation]
+                        cond = [
+                            pred[:, class_idx]
+                            if self.consecutive_slices == 1 or self.dimensionality == 2
+                            else pred[:, :, class_idx]
+                            for pred in predictions_segmentation
+                        ]
                     else:
-                        cond = predictions_segmentation[:, class_idx]
+                        cond = (
+                            predictions_segmentation[:, class_idx]
+                            if (self.consecutive_slices == 1 or self.dimensionality == 2)
+                            else predictions_segmentation[:, :, class_idx]
+                        )
 
                 if isinstance(predictions_segmentation, list):
                     for idx, pred in enumerate(predictions_segmentation):
-                        predictions_segmentation[idx][:, class_idx] = torch.where(
-                            cond[idx] >= thres,
-                            predictions_segmentation[idx][:, class_idx],
-                            torch.zeros_like(predictions_segmentation[idx][:, class_idx]),
-                        )
+                        if self.consecutive_slices == 1 or self.dimensionality == 2:
+                            predictions_segmentation[idx][:, class_idx] = torch.where(
+                                cond[idx] >= thres,
+                                predictions_segmentation[idx][:, class_idx],
+                                torch.zeros_like(predictions_segmentation[idx][:, class_idx]),
+                            )
+                        else:
+                            predictions_segmentation[idx][:, :, class_idx] = torch.where(
+                                cond[idx] >= thres,
+                                predictions_segmentation[idx][:, :, class_idx],
+                                torch.zeros_like(predictions_segmentation[idx][:, :, class_idx]),
+                            )
                 else:
-                    predictions_segmentation[:, class_idx] = torch.where(
-                        cond >= thres,
-                        predictions_segmentation[:, class_idx],
-                        torch.zeros_like(predictions_segmentation[:, class_idx]),
-                    )
+                    if self.consecutive_slices == 1 or self.dimensionality == 2:
+                        predictions_segmentation[:, class_idx] = torch.where(
+                            cond >= thres,
+                            predictions_segmentation[:, class_idx],
+                            torch.zeros_like(predictions_segmentation[:, class_idx]),
+                        )
+                    else:
+                        predictions_segmentation[:, :, class_idx] = torch.where(
+                            cond >= thres,
+                            predictions_segmentation[:, :, class_idx],
+                            torch.zeros_like(predictions_segmentation[:, :, class_idx]),
+                        )
 
         # Noise-to-Recon forward pass, if Noise-to-Recon is used.
         predictions_reconstruction_n2r = None
@@ -1310,11 +1334,13 @@ def inference_step(  # noqa: MC0001
         # Get acceleration factor from acceleration list, if multiple accelerations are used. Or if batch size > 1.
         if isinstance(acceleration, list):
             if acceleration[0].shape[0] > 1:
-                acceleration[0] = acceleration[0][0]
+                for i in enumerate(acceleration):
+                    acceleration[i] = acceleration[i][0]
             acceleration = np.round(acceleration[r].item())
         else:
-            if acceleration.shape[0] > 1:  # type: ignore
-                acceleration = acceleration[0]  # type: ignore
+            if acceleration[0].shape[0] > 1:  # type: ignore
+                for i in enumerate(acceleration):  # type: ignore
+                    acceleration[i] = acceleration[i][0]  # type: ignore
             acceleration = np.round(acceleration.item())  # type: ignore
 
         # Pass r to the attrs dictionary, so that it can be used in unnormalize_for_loss_or_log if needed.

atommic/collections/multitask/rs/nn/mtlrs.py

@@ -160,6 +160,8 @@ def forward(
             init_reconstruction_pred = pred_reconstruction[-1][-1]
 
             if self.task_adaption_type == "multi_task_learning":
+                # this is the most important step of the mtlrs model, where the segmentation mask is applied to the
+                # hidden states and updates them
                 hidden_states = [
                     torch.cat(
                         [torch.abs(init_reconstruction_pred.unsqueeze(self.coil_dim) * pred_segmentation)]
@@ -170,21 +172,17 @@ def forward(
                     if f != 0
                 ]
 
-                if self.consecutive_slices > 1:
-                    hx = [x.unsqueeze(1) for x in hx]
-
                 # Check if the concatenated hidden states are the same size as the hidden state of the RNN
-                if hidden_states[0].shape[self.coil_dim] != hx[0].shape[self.coil_dim]:
-                    prev_hidden_states = hidden_states
-                    hidden_states = []
-                    for hs in prev_hidden_states:
-                        new_hidden_state = hs
-                        for _ in range(hx[0].shape[1] - prev_hidden_states[0].shape[1]):
-                            new_hidden_state = torch.cat(
-                                [new_hidden_state, torch.zeros_like(hx[0][:, 0, :, :]).unsqueeze(self.coil_dim)],
-                                dim=self.coil_dim,
-                            )
-                        hidden_states.append(new_hidden_state)
+                if (
+                    hidden_states[0].shape[self.coil_dim if self.consecutive_slices == 1 else self.coil_dim - 1]
+                    != hx[0].shape[self.coil_dim if self.consecutive_slices == 1 else self.coil_dim - 1]
+                ):
+                    hidden_states = [
+                        hidden_states[i].reshape(
+                            hidden_states[0].shape[0] * self.consecutive_slices, *hidden_states[i].shape[2:]
+                        )
+                        for i in range(len(hidden_states))
+                    ]
 
                 hx = [hx[i] + hidden_states[i] for i in range(len(hx))]
 

atommic/collections/multitask/rs/nn/mtlrs_base/mtlrs_block.py

@@ -82,8 +82,6 @@ def __init__(
         self.spatial_dims = spatial_dims
         self.coil_dim = coil_dim
         self.dimensionality = dimensionality
-        if self.dimensionality != 2:
-            raise NotImplementedError(f"Currently only 2D is supported for segmentation, got {self.dimensionality}D.")
         self.consecutive_slices = consecutive_slices
         self.coil_combination_method = coil_combination_method
 
@@ -237,7 +235,6 @@ def forward(  # noqa: MC0001
                 )
                 cascades_predictions = []
                 for i, cascade in enumerate(self.reconstruction_module):
-                    # Forward pass through the cascades
                     prediction_slice, hx = cascade(
                         prediction_slice,
                         y_slice,
@@ -248,6 +245,11 @@ def forward(  # noqa: MC0001
                         sigma,
                         keep_prediction=False if i == 0 else self.keep_prediction,
                     )
+                    if (prediction_slice[0].shape[0] == self.consecutive_slices) and (
+                        hx[0].shape[0] == self.consecutive_slices
+                    ):
+                        prediction_slice = [pred[slice_idx].unsqueeze(0) for pred in prediction_slice]
+                        hx = [h[slice_idx].unsqueeze(0) for h in hx]
                     time_steps_predictions = [torch.view_as_complex(pred) for pred in prediction_slice]
                     cascades_predictions.append(torch.stack(time_steps_predictions, dim=0))
                 pred_reconstruction_slices.append(torch.stack(cascades_predictions, dim=0))
@@ -267,6 +269,8 @@ def forward(  # noqa: MC0001
                 if init_reconstruction_pred is None or init_reconstruction_pred.dim() < 4
                 else init_reconstruction_pred
             )
+            if self.consecutive_slices > 1 and self.reconstruction_module_dimensionality == 3:
+                mask = torch.concatenate([mask.unsqueeze(self.coil_dim)] * y.shape[1], 1)
             sigma = 1.0
             cascades_predictions = []
             for i, cascade in enumerate(self.reconstruction_module):
@@ -292,11 +296,13 @@ def forward(  # noqa: MC0001
             _pred_reconstruction = _pred_reconstruction[-1]
         if _pred_reconstruction.shape[-1] != 2:
             _pred_reconstruction = torch.view_as_real(_pred_reconstruction)
+
         if self.consecutive_slices > 1 and _pred_reconstruction.dim() == 5:
             _pred_reconstruction = _pred_reconstruction.reshape(
                 _pred_reconstruction.shape[0] * _pred_reconstruction.shape[1],
                 *_pred_reconstruction.shape[2:],
             )
+
         if _pred_reconstruction.shape[-1] == 2:
             if self.input_channels == 1:
                 _pred_reconstruction = torch.view_as_complex(_pred_reconstruction).unsqueeze(1)