diff --git a/configs/callbacks/combined_test_suite_mil.yaml b/configs/callbacks/combined_test_suite_mil.yaml new file mode 100644 index 00000000..4efd9971 --- /dev/null +++ b/configs/callbacks/combined_test_suite_mil.yaml @@ -0,0 +1,6 @@ +defaults: + - curves_mil + - nested_metrics_mil + - tile_histograms_mil + - mil_prediction_callback + - _self_ diff --git a/configs/callbacks/curves.yaml b/configs/callbacks/curves.yaml index d848857b..7cd9d2dc 100644 --- a/configs/callbacks/curves.yaml +++ b/configs/callbacks/curves.yaml @@ -1,4 +1,4 @@ curves_callback: - _target_: prostate_cancer.callbacks.CurvesCallback + _target_: prostate_cancer.callbacks.CurvesCallbackTile threshold: ${tile_threshold} optimal_seek: false diff --git a/configs/callbacks/curves_mil.yaml b/configs/callbacks/curves_mil.yaml new file mode 100644 index 00000000..ee55fe20 --- /dev/null +++ b/configs/callbacks/curves_mil.yaml @@ -0,0 +1,4 @@ +curves_callback: + _target_: prostate_cancer.callbacks.CurvesCallbackMIL + threshold: ${tile_threshold} + optimal_seek: false diff --git a/configs/callbacks/curves_sl.yaml b/configs/callbacks/curves_sl.yaml new file mode 100644 index 00000000..4604407f --- /dev/null +++ b/configs/callbacks/curves_sl.yaml @@ -0,0 +1,4 @@ +curves_callback: + _target_: prostate_cancer.callbacks.CurvesCallbackSL + threshold: ${slide_threshold} + optimal_seek: false diff --git a/configs/callbacks/nested_metrics_mil.yaml b/configs/callbacks/nested_metrics_mil.yaml new file mode 100644 index 00000000..a340547c --- /dev/null +++ b/configs/callbacks/nested_metrics_mil.yaml @@ -0,0 +1,3 @@ +nested_metrics_callback: + _target_: prostate_cancer.callbacks.NestedMetricsCallbackMIL + threshold: ${tile_threshold} diff --git a/configs/callbacks/slide_histograms_mil.yaml b/configs/callbacks/slide_histograms_mil.yaml new file mode 100644 index 00000000..6730cd43 --- /dev/null +++ b/configs/callbacks/slide_histograms_mil.yaml @@ -0,0 +1,2 @@ +slide_histograms_callback: + _target_: prostate_cancer.callbacks.SlideHistogramsCallbackMIL diff --git a/configs/callbacks/slide_threshold_mil.yaml b/configs/callbacks/slide_threshold_mil.yaml new file mode 100644 index 00000000..f442cf3b --- /dev/null +++ b/configs/callbacks/slide_threshold_mil.yaml @@ -0,0 +1,7 @@ +defaults: + - slide_histograms_mil + - curves_sl + - _self_ + +curves_callback: + optimal_seek: true diff --git a/configs/callbacks/tile_histograms.yaml b/configs/callbacks/tile_histograms.yaml index 16a4bfd0..8f1e765b 100644 --- a/configs/callbacks/tile_histograms.yaml +++ b/configs/callbacks/tile_histograms.yaml @@ -1,2 +1,2 @@ tile_histograms_callback: - _target_: prostate_cancer.callbacks.TileHistogramsCallback + _target_: prostate_cancer.callbacks.TileHistogramsCallbackTile diff --git a/configs/callbacks/tile_histograms_mil.yaml b/configs/callbacks/tile_histograms_mil.yaml new file mode 100644 index 00000000..87f8a110 --- /dev/null +++ b/configs/callbacks/tile_histograms_mil.yaml @@ -0,0 +1,2 @@ +tile_histograms_callback: + _target_: prostate_cancer.callbacks.TileHistogramsCallbackMIL diff --git a/configs/callbacks/tile_threshold_mil.yaml b/configs/callbacks/tile_threshold_mil.yaml new file mode 100644 index 00000000..5ac964c4 --- /dev/null +++ b/configs/callbacks/tile_threshold_mil.yaml @@ -0,0 +1,7 @@ +defaults: + - tile_histograms_mil + - curves_mil + - _self_ + +curves_callback: + optimal_seek: true diff --git a/configs/experiment/predict/bag_of_embeddings/pgp/base.yaml b/configs/experiment/predict/bag_of_embeddings/pgp/base.yaml index 318d8b41..78f975da 100644 --- a/configs/experiment/predict/bag_of_embeddings/pgp/base.yaml +++ b/configs/experiment/predict/bag_of_embeddings/pgp/base.yaml @@ -2,8 +2,8 @@ defaults: - /experiment/predict/bag_of_embeddings/base - - /model/pgp_mil_model@model - - /model/checkpoints/pgp_mil@model_data + - /model/pgp_hybrid_mil_model@model + - /model/checkpoints/mmci_tl/pgp_hybrid_mil@model_data - _self_ datamodule: diff --git a/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/base.yaml b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/base.yaml index 35419b5f..3d239b1d 100644 --- a/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/base.yaml +++ b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/base.yaml @@ -1,7 +1,7 @@ # @package _global_ defaults: - - /callbacks/mil_predictions@callbacks + - /callbacks/mil_prediction_callback@callbacks - /experiment/predict/bag_of_embeddings/pgp/base - _self_ diff --git a/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_sl_group_1.yaml b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_sl_group_1.yaml new file mode 100644 index 00000000..ad7964be --- /dev/null +++ b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_sl_group_1.yaml @@ -0,0 +1,6 @@ +# @package _global_ + +defaults: + - /experiment/predict/bag_of_embeddings/pgp/mil_predictions/base + - /data/mmci/sl_group_1@predict_data + - _self_ diff --git a/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_tl_val.yaml b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_tl_val.yaml index 353cf3c1..c0c2aeba 100644 --- a/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_tl_val.yaml +++ b/configs/experiment/predict/bag_of_embeddings/pgp/mil_predictions/mmci_tl_val.yaml @@ -2,5 +2,5 @@ defaults: - /experiment/predict/bag_of_embeddings/pgp/mil_predictions/base - - /data/mmci_tl_val@predict_data + - /data/mmci/tl_val@predict_data - _self_ diff --git a/configs/experiment/predict/bag_of_embeddings/pgp/slide_threshold.yaml b/configs/experiment/predict/bag_of_embeddings/pgp/slide_threshold.yaml new file mode 100644 index 00000000..11f1d550 --- /dev/null +++ b/configs/experiment/predict/bag_of_embeddings/pgp/slide_threshold.yaml @@ -0,0 +1,10 @@ +# @package _global_ + +defaults: + - /callbacks/slide_threshold_mil@callbacks + - /experiment/predict/bag_of_embeddings/pgp/base + - /data/mmci/sl_group_1@predict_data + - _self_ + +slide_threshold: 0.99 # My estimate for other models +use_case: "SL Threshold Estimation" diff --git a/configs/experiment/predict/bag_of_embeddings/virchow2/base.yaml b/configs/experiment/predict/bag_of_embeddings/virchow2/base.yaml index 0748a564..cff0ca13 100644 --- a/configs/experiment/predict/bag_of_embeddings/virchow2/base.yaml +++ b/configs/experiment/predict/bag_of_embeddings/virchow2/base.yaml @@ -2,8 +2,8 @@ defaults: - /experiment/predict/bag_of_embeddings/base - - /model/virchow2_mil_model@model - - /model/checkpoints/mmci_tl/virchow2_mil@model_data + - /model/virchow2_hybrid_mil_model@model + - /model/checkpoints/mmci_tl/virchow2_hybrid_mil@model_data - _self_ datamodule: diff --git a/configs/experiment/predict/bag_of_embeddings/virchow2/mil_predictions/mmci_sl_group_1.yaml b/configs/experiment/predict/bag_of_embeddings/virchow2/mil_predictions/mmci_sl_group_1.yaml new file mode 100644 index 00000000..fe9179fb --- /dev/null +++ b/configs/experiment/predict/bag_of_embeddings/virchow2/mil_predictions/mmci_sl_group_1.yaml @@ -0,0 +1,6 @@ +# @package _global_ + +defaults: + - /experiment/predict/bag_of_embeddings/virchow2/mil_predictions/base + - /data/mmci/sl_group_1@predict_data + - _self_ diff --git a/configs/experiment/predict/bag_of_embeddings/virchow2/slide_threshold.yaml b/configs/experiment/predict/bag_of_embeddings/virchow2/slide_threshold.yaml new file mode 100644 index 00000000..dbcd62b5 --- /dev/null +++ b/configs/experiment/predict/bag_of_embeddings/virchow2/slide_threshold.yaml @@ -0,0 +1,10 @@ +# @package _global_ + +defaults: + - /callbacks/slide_threshold_mil@callbacks + - /experiment/predict/bag_of_embeddings/virchow2/base + - /data/mmci/sl_group_1@predict_data + - _self_ + +slide_threshold: 0.99 # My estimate for other models +use_case: "SL Threshold Estimation" diff --git a/configs/experiment/test/bag_of_embeddings/base.yaml b/configs/experiment/test/bag_of_embeddings/base.yaml new file mode 100644 index 00000000..919deb91 --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/base.yaml @@ -0,0 +1,14 @@ +# @package _global_ + +defaults: + - /experiment/test/base + - /datamodule/bag_of_embeddings_datamodule@datamodule + - _self_ + +slide_threshold: ${model_data.slide_threshold} +model: + sl_threshold: ${slide_threshold} + +datamodule: + batch_size: 4 + num_workers: 8 diff --git a/configs/experiment/test/bag_of_embeddings/pgp/base.yaml b/configs/experiment/test/bag_of_embeddings/pgp/base.yaml new file mode 100644 index 00000000..25db044e --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/pgp/base.yaml @@ -0,0 +1,12 @@ +# @package _global_ + +defaults: + - /experiment/test/bag_of_embeddings/base + - /model/pgp_hybrid_mil_model@model + - /model/checkpoints/mmci_tl/pgp_hybrid_mil@model_data + - _self_ + +datamodule: + test: + uris: + - ${test_data.tiles_filtered_w_pgp_uri_224} diff --git a/configs/experiment/test/bag_of_embeddings/pgp/test_suite/base.yaml b/configs/experiment/test/bag_of_embeddings/pgp/test_suite/base.yaml new file mode 100644 index 00000000..32e8817b --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/pgp/test_suite/base.yaml @@ -0,0 +1,8 @@ +# @package _global_ + +defaults: + - /callbacks/combined_test_suite_mil@callbacks + - /experiment/test/bag_of_embeddings/pgp/base + - _self_ + +use_case: "TL Test" diff --git a/configs/experiment/test/bag_of_embeddings/pgp/test_suite/mmci_tl_test.yaml b/configs/experiment/test/bag_of_embeddings/pgp/test_suite/mmci_tl_test.yaml new file mode 100644 index 00000000..f5f43efc --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/pgp/test_suite/mmci_tl_test.yaml @@ -0,0 +1,6 @@ +# @package _global_ + +defaults: + - /experiment/test/bag_of_embeddings/pgp/test_suite/base + - /data/mmci/tl_test@test_data + - _self_ diff --git a/configs/experiment/test/bag_of_embeddings/pgp/tile_threshold.yaml b/configs/experiment/test/bag_of_embeddings/pgp/tile_threshold.yaml new file mode 100644 index 00000000..1e133ddd --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/pgp/tile_threshold.yaml @@ -0,0 +1,10 @@ +# @package _global_ + +defaults: + - /callbacks/tile_threshold_mil@callbacks + - /experiment/test/bag_of_embeddings/pgp/base + - /data/mmci/tl_val@test_data + - _self_ + +tile_threshold: 0.9 # Pathologist's estimate for other models +use_case: "TL Threshold Estimation" diff --git a/configs/experiment/test/bag_of_embeddings/virchow2/base.yaml b/configs/experiment/test/bag_of_embeddings/virchow2/base.yaml new file mode 100644 index 00000000..c90e5378 --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/virchow2/base.yaml @@ -0,0 +1,12 @@ +# @package _global_ + +defaults: + - /experiment/test/bag_of_embeddings/base + - /model/virchow2_hybrid_mil_model@model + - /model/checkpoints/mmci_tl/virchow2_hybrid_mil@model_data + - _self_ + +datamodule: + test: + uris: + - ${test_data.tiles_filtered_w_virchow2_uri_224} diff --git a/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/base.yaml b/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/base.yaml new file mode 100644 index 00000000..2bc6d4e4 --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/base.yaml @@ -0,0 +1,8 @@ +# @package _global_ + +defaults: + - /callbacks/combined_test_suite_mil@callbacks + - /experiment/test/bag_of_embeddings/virchow2/base + - _self_ + +use_case: "TL Test" diff --git a/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/mmci_tl_test.yaml b/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/mmci_tl_test.yaml new file mode 100644 index 00000000..6bb9ae18 --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/virchow2/test_suite/mmci_tl_test.yaml @@ -0,0 +1,6 @@ +# @package _global_ + +defaults: + - /experiment/test/bag_of_embeddings/virchow2/test_suite/base + - /data/mmci/tl_test@test_data + - _self_ diff --git a/configs/experiment/test/bag_of_embeddings/virchow2/tile_threshold.yaml b/configs/experiment/test/bag_of_embeddings/virchow2/tile_threshold.yaml new file mode 100644 index 00000000..8d5ac728 --- /dev/null +++ b/configs/experiment/test/bag_of_embeddings/virchow2/tile_threshold.yaml @@ -0,0 +1,10 @@ +# @package _global_ + +defaults: + - /callbacks/tile_threshold_mil@callbacks + - /experiment/test/bag_of_embeddings/virchow2/base + - /data/mmci/tl_val@test_data + - _self_ + +tile_threshold: 0.9 # Pathologist's estimate for other models +use_case: "TL Threshold Estimation" diff --git a/configs/experiment/train/bag_of_embeddings/pgp.yaml b/configs/experiment/train/bag_of_embeddings/pgp.yaml index a3be9998..8c4439d0 100644 --- a/configs/experiment/train/bag_of_embeddings/pgp.yaml +++ b/configs/experiment/train/bag_of_embeddings/pgp.yaml @@ -2,7 +2,7 @@ defaults: - /experiment/train/bag_of_embeddings/base - - /model/pgp_mil_model@model + - /model/pgp_hybrid_mil_model@model - _self_ datamodule: diff --git a/configs/experiment/train/bag_of_embeddings/virchow2.yaml b/configs/experiment/train/bag_of_embeddings/virchow2.yaml index 4a3674e8..15adeb39 100644 --- a/configs/experiment/train/bag_of_embeddings/virchow2.yaml +++ b/configs/experiment/train/bag_of_embeddings/virchow2.yaml @@ -2,7 +2,7 @@ defaults: - /experiment/train/bag_of_embeddings/base - - /model/virchow2_mil_model@model + - /model/virchow2_hybrid_mil_model@model - _self_ datamodule: diff --git a/configs/model/checkpoints/mmci_tl/pgp_mil.yaml b/configs/model/checkpoints/mmci_tl/pgp_hybrid_mil.yaml similarity index 79% rename from configs/model/checkpoints/mmci_tl/pgp_mil.yaml rename to configs/model/checkpoints/mmci_tl/pgp_hybrid_mil.yaml index 074c4921..4c5b99d9 100644 --- a/configs/model/checkpoints/mmci_tl/pgp_mil.yaml +++ b/configs/model/checkpoints/mmci_tl/pgp_hybrid_mil.yaml @@ -2,8 +2,8 @@ defaults: - base_foundation - _self_ -tile_threshold: 0.5 -slide_threshold: 0.5 +tile_threshold: 0.88 +slide_threshold: 0.9783778786 checkpoint: "mlflow-artifacts:/65/7cb992adfe4b422790becc5984ff2a6f/artifacts/checkpoints/epoch=8-step=1583/checkpoint.ckpt" model_name: "PGP (Hybrid-MIL)" diff --git a/configs/model/checkpoints/mmci_tl/virchow2_mil.yaml b/configs/model/checkpoints/mmci_tl/virchow2_hybrid_mil.yaml similarity index 79% rename from configs/model/checkpoints/mmci_tl/virchow2_mil.yaml rename to configs/model/checkpoints/mmci_tl/virchow2_hybrid_mil.yaml index 9a00458b..11f36170 100644 --- a/configs/model/checkpoints/mmci_tl/virchow2_mil.yaml +++ b/configs/model/checkpoints/mmci_tl/virchow2_hybrid_mil.yaml @@ -2,8 +2,8 @@ defaults: - base_foundation - _self_ -tile_threshold: 0.5 -slide_threshold: 0.5 +tile_threshold: 0.85 +slide_threshold: 0.9917988777 checkpoint: "mlflow-artifacts:/65/5e0c1d4699ca49c4938baf3c65e7f61c/artifacts/checkpoints/epoch=6-step=1176/checkpoint.ckpt" model_name: "Virchow2 (Hybrid-MIL)" diff --git a/configs/model/pgp_mil_model.yaml b/configs/model/pgp_hybrid_mil_model.yaml similarity index 100% rename from configs/model/pgp_mil_model.yaml rename to configs/model/pgp_hybrid_mil_model.yaml diff --git a/configs/model/virchow2_mil_model.yaml b/configs/model/virchow2_hybrid_mil_model.yaml similarity index 100% rename from configs/model/virchow2_mil_model.yaml rename to configs/model/virchow2_hybrid_mil_model.yaml diff --git a/prostate_cancer/attention_mil_model.py b/prostate_cancer/attention_mil_model.py index e385234b..dd424281 100644 --- a/prostate_cancer/attention_mil_model.py +++ b/prostate_cancer/attention_mil_model.py @@ -203,10 +203,10 @@ def validation_step(self, batch: LabeledBagOfTilesSampleBatch) -> None: self.val_metrics_tl, on_epoch=True, on_step=False, batch_size=len(bags) ) - def test_step(self, batch: LabeledBagOfTilesSampleBatch) -> None: + def test_step(self, batch: LabeledBagOfTilesSampleBatch) -> MILModelOutput: # type: ignore[override] bags, tl_labels, sl_labels, _ = batch - sl_outputs, tl_outputs, mask, _ = self(bags) + sl_outputs, tl_outputs, mask, attention = self(bags) self.test_metrics_sl.update(sl_outputs, sl_labels) self.test_metrics_tl.update(tl_outputs[mask.bool()], tl_labels[mask.bool()]) @@ -217,6 +217,7 @@ def test_step(self, batch: LabeledBagOfTilesSampleBatch) -> None: self.log_dict( self.test_metrics_tl, on_epoch=True, on_step=False, batch_size=len(bags) ) + return sl_outputs.sigmoid(), tl_outputs.sigmoid(), mask, attention def predict_step(self, batch: UnlabeledBagOfTilesSampleBatch) -> MILModelOutput: sl_preds_raw, tl_preds_raw, mask, attention = self(batch[0]) diff --git a/prostate_cancer/callbacks/__init__.py b/prostate_cancer/callbacks/__init__.py index 60089480..875e6156 100644 --- a/prostate_cancer/callbacks/__init__.py +++ b/prostate_cancer/callbacks/__init__.py @@ -3,7 +3,9 @@ from prostate_cancer.callbacks.carcinoma_prediction_table_callback import ( CarcinomaPredictionTableCallback, ) -from prostate_cancer.callbacks.curves_callback import CurvesCallback +from prostate_cancer.callbacks.curves_callback_mil import CurvesCallbackMIL +from prostate_cancer.callbacks.curves_callback_sl import CurvesCallbackSL +from prostate_cancer.callbacks.curves_callback_tl import CurvesCallbackTile from prostate_cancer.callbacks.estimation_callback import ( EstimationCallback, ) @@ -13,20 +15,36 @@ MultiAggregatorEvalCallback, ) from prostate_cancer.callbacks.nested_metrics_callback import NestedMetricsCallback +from prostate_cancer.callbacks.nested_metrics_callback_mil import ( + NestedMetricsCallbackMIL, +) from prostate_cancer.callbacks.num_positive_callback import NumPositiveCallback -from prostate_cancer.callbacks.tile_histograms_callback import TileHistogramsCallback +from prostate_cancer.callbacks.slide_histograms_callback_mil import ( + SlideHistogramsCallbackMIL, +) +from prostate_cancer.callbacks.tile_histograms_callback_mil import ( + TileHistogramsCallbackMIL, +) +from prostate_cancer.callbacks.tile_histograms_callback_tl import ( + TileHistogramsCallbackTile, +) __all__ = [ "AggregatorCallback", "CAMExplainer", "CarcinomaPredictionTableCallback", - "CurvesCallback", + "CurvesCallbackMIL", + "CurvesCallbackSL", + "CurvesCallbackTile", "EstimationCallback", "HeatmapCallback", "MILPredictionCallback", "MultiAggregatorEvalCallback", "NestedMetricsCallback", + "NestedMetricsCallbackMIL", "NumPositiveCallback", - "TileHistogramsCallback", + "SlideHistogramsCallbackMIL", + "TileHistogramsCallbackMIL", + "TileHistogramsCallbackTile", ] diff --git a/prostate_cancer/callbacks/curves_callback.py b/prostate_cancer/callbacks/curves_callback_base.py similarity index 78% rename from prostate_cancer/callbacks/curves_callback.py rename to prostate_cancer/callbacks/curves_callback_base.py index d5d58649..77155eb2 100644 --- a/prostate_cancer/callbacks/curves_callback.py +++ b/prostate_cancer/callbacks/curves_callback_base.py @@ -1,4 +1,4 @@ -from typing import Any +from abc import ABC import mlflow import numpy as np @@ -8,36 +8,26 @@ from sklearn.metrics import auc, precision_recall_curve, roc_curve from postprocessing.slide_level_curves import _plot_curve -from prostate_cancer.typing import LabeledTileSampleBatch -class CurvesCallback(Callback): - def __init__(self, threshold: float, optimal_seek: bool = True) -> None: - """This callback creates tile-level ROC curve and Precision-Recall curve and marks selected + optimized thresholds used for metric computation. +class CurvesCallbackBase(Callback, ABC): + def __init__( + self, threshold: float, tile_level: bool, optimal_seek: bool = True + ) -> None: + """This callback creates tile-level or slide-level ROC curve and Precision-Recall curve and marks selected + optimized thresholds used for metric computation. Args: threshold (float): pathologist selected threshold + tile_level (bool): whether the curves are computed on tile-level or slide-level predictions optimal_seek (bool): whether we are looking for optimal thresholds or just want to plot the curves """ super().__init__() + self.tile_level = tile_level self.optimal_seek = optimal_seek self.threshold = threshold self.preds: list[torch.Tensor] = [] self.targets: list[torch.Tensor] = [] - def on_test_batch_end( - self, - trainer: Trainer, - pl_module: LightningModule, - outputs: Any, - batch: LabeledTileSampleBatch, - batch_idx: int, - dataloader_idx: int = 0, - ) -> None: - targets = batch[1] - self.preds.append(outputs.cpu()) - self.targets.append(targets.cpu()) - def _plot_roc( self, y_pred: NDArray[np.float32], y_true: NDArray[np.float32] ) -> None: @@ -59,6 +49,7 @@ def _plot_roc( j = tpr - fpr optimal_idx = j.argmax() j_threshold = roc_thresholds[optimal_idx] + mlflow.log_param("j_threshold", j_threshold) j_fpr = fpr[optimal_idx] j_tpr = tpr[optimal_idx] @@ -66,7 +57,7 @@ def _plot_roc( labels.append(f"J Threshold = {j_threshold:.2f}") colors.append("green") - plot_path = "tile_roc.png" + plot_path = "tile_roc.png" if self.tile_level else "slide_roc.png" _plot_curve( fpr, tpr, @@ -99,12 +90,17 @@ def _plot_precision_recall( f1 = 2 * (precision * recall) / (precision + recall + 1e-8) best_idx = np.argmax(f1) best_threshold = thresholds[best_idx] + mlflow.log_param("pr_threshold", best_threshold) to_pinpoint.append((recall[best_idx], precision[best_idx])) labels.append(f"F1 Threshold = {best_threshold:.2f}") colors.append("green") - plot_path = "tile_precision_recall.png" + plot_path = ( + "tile_precision_recall.png" + if self.tile_level + else "slide_precision_recall.png" + ) _plot_curve( recall, precision, @@ -115,12 +111,12 @@ def _plot_precision_recall( "Recall", "Precision", "Precision-Recall Curve", - "tile_precision_recall.png", + plot_path, "lower left", ) mlflow.log_artifact(plot_path, artifact_path="plots") - def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + def _plot_and_clear(self) -> None: y_pred = torch.cat(self.preds).numpy() y_true = torch.cat(self.targets).numpy() @@ -129,3 +125,11 @@ def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> Non self.preds.clear() self.targets.clear() + + def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + self._plot_and_clear() + + def on_predict_epoch_end( + self, trainer: Trainer, pl_module: LightningModule + ) -> None: + self._plot_and_clear() diff --git a/prostate_cancer/callbacks/curves_callback_mil.py b/prostate_cancer/callbacks/curves_callback_mil.py new file mode 100644 index 00000000..ae3ad927 --- /dev/null +++ b/prostate_cancer/callbacks/curves_callback_mil.py @@ -0,0 +1,27 @@ +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.curves_callback_base import CurvesCallbackBase +from prostate_cancer.typing import LabeledBagOfTilesSampleBatch, MILModelOutput + + +class CurvesCallbackMIL(CurvesCallbackBase): + def __init__(self, threshold: float, optimal_seek: bool) -> None: + super().__init__( + threshold=threshold, tile_level=True, optimal_seek=optimal_seek + ) + + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, # type: ignore[override] + batch: LabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + + _, tl_outputs_raw, mask, _ = outputs + tl_outputs_valid = tl_outputs_raw[mask.bool()] + targets = batch[1][mask.bool()] + self.preds.append(tl_outputs_valid.cpu()) + self.targets.append(targets.cpu()) diff --git a/prostate_cancer/callbacks/curves_callback_sl.py b/prostate_cancer/callbacks/curves_callback_sl.py new file mode 100644 index 00000000..9a2bca89 --- /dev/null +++ b/prostate_cancer/callbacks/curves_callback_sl.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING, cast + +import torch +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.curves_callback_base import CurvesCallbackBase +from prostate_cancer.typing import MILModelOutput, UnlabeledBagOfTilesSampleBatch + + +if TYPE_CHECKING: + from prostate_cancer.datamodule import BagOfTilesDataModule + + +class CurvesCallbackSL(CurvesCallbackBase): + def __init__(self, threshold: float, optimal_seek: bool) -> None: + super().__init__( + threshold=threshold, tile_level=False, optimal_seek=optimal_seek + ) + + def setup( + self, trainer: Trainer, pl_module: LightningModule, stage: str | None = None + ) -> None: + if not hasattr(trainer, "datamodule"): + raise ValueError("Trainer should have datamodule attribute") + + datamodule = cast("BagOfTilesDataModule", trainer.datamodule) + slides = datamodule.predict.slides + self._slide_targets = dict(zip(slides["id"], slides["carcinoma"], strict=True)) + + def on_predict_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, + batch: UnlabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + sl_outputs, _, _, _ = outputs + _, metadata_batch = batch + + targets = torch.tensor( + [float(self._slide_targets[m["slide_id"]]) for m in metadata_batch] + ) + + self.preds.append(sl_outputs.detach().cpu()) + self.targets.append(targets) diff --git a/prostate_cancer/callbacks/curves_callback_tl.py b/prostate_cancer/callbacks/curves_callback_tl.py new file mode 100644 index 00000000..51b31bb3 --- /dev/null +++ b/prostate_cancer/callbacks/curves_callback_tl.py @@ -0,0 +1,26 @@ +from typing import Any + +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.curves_callback_base import CurvesCallbackBase +from prostate_cancer.typing import LabeledTileSampleBatch + + +class CurvesCallbackTile(CurvesCallbackBase): + def __init__(self, threshold: float, optimal_seek: bool) -> None: + super().__init__( + threshold=threshold, tile_level=True, optimal_seek=optimal_seek + ) + + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: Any, + batch: LabeledTileSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + targets = batch[1] + self.preds.append(outputs.cpu()) + self.targets.append(targets.cpu()) diff --git a/prostate_cancer/callbacks/tile_histograms_callback.py b/prostate_cancer/callbacks/histograms_callback_base.py similarity index 76% rename from prostate_cancer/callbacks/tile_histograms_callback.py rename to prostate_cancer/callbacks/histograms_callback_base.py index ca304609..a50f4fe7 100644 --- a/prostate_cancer/callbacks/tile_histograms_callback.py +++ b/prostate_cancer/callbacks/histograms_callback_base.py @@ -1,5 +1,5 @@ +from abc import ABC from pathlib import Path -from typing import Any import matplotlib.pyplot as plt import mlflow @@ -8,33 +8,15 @@ from numpy.typing import NDArray from rationai.mlkit.lightning.loggers import MLFlowLogger -from prostate_cancer.typing import LabeledTileSampleBatch - -class TileHistogramsCallback(Callback): +class HistogramsCallbackBase(Callback, ABC): def __init__(self) -> None: """This callback creates prediction histograms for both negative and positive distribution of tiles.""" super().__init__() self.all_preds: list[NDArray[np.floating]] = [] self.all_labels: list[NDArray[np.floating]] = [] - def on_test_batch_end( - self, - trainer: Trainer, - pl_module: LightningModule, - outputs: Any, - batch: LabeledTileSampleBatch, - batch_idx: int, - dataloader_idx: int = 0, - ) -> None: - _, y, _ = batch - preds = outputs.detach().cpu().numpy().flatten() - labels = y.detach().cpu().numpy().flatten() - - self.all_preds.append(preds) - self.all_labels.append(labels) - - def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + def _plot_and_clear(self, trainer: Trainer) -> None: assert isinstance(trainer.logger, MLFlowLogger) preds = np.concatenate(self.all_preds) @@ -68,3 +50,11 @@ def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> Non self.all_preds.clear() self.all_labels.clear() + + def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + self._plot_and_clear(trainer) + + def on_predict_epoch_end( + self, trainer: Trainer, pl_module: LightningModule + ) -> None: + self._plot_and_clear(trainer) diff --git a/prostate_cancer/callbacks/mil_prediction_callback.py b/prostate_cancer/callbacks/mil_prediction_callback.py index 2a642545..9e847262 100644 --- a/prostate_cancer/callbacks/mil_prediction_callback.py +++ b/prostate_cancer/callbacks/mil_prediction_callback.py @@ -9,7 +9,11 @@ from rationai.masks.mask_builders import ScalarMaskBuilder from rationai.mlkit.lightning.loggers.mlflow import MLFlowLogger -from prostate_cancer.typing import MILModelOutput, UnlabeledBagOfTilesSampleBatch +from prostate_cancer.typing import ( + LabeledBagOfTilesSampleBatch, + MILModelOutput, + UnlabeledBagOfTilesSampleBatch, +) if TYPE_CHECKING: @@ -28,7 +32,8 @@ def setup( raise ValueError("Trainer should have datamodule attribute") datamodule = cast("BagOfTilesDataModule", trainer.datamodule) - slides = cast("HFDataset", datamodule.predict.slides) + dataset = datamodule.test if stage == "test" else datamodule.predict + slides = cast("HFDataset", dataset.slides) self._slide_index = { Path(path).stem: i for i, path in enumerate(slides["path"]) @@ -67,20 +72,16 @@ def get_mask_builder( stride=slide["stride_x"], ) - def on_predict_batch_end( + def _on_batch_end( self, trainer: Trainer, - pl_module: LightningModule, outputs: MILModelOutput, - batch: UnlabeledBagOfTilesSampleBatch, - batch_idx: int, - dataloader_idx: int = 0, + batch: UnlabeledBagOfTilesSampleBatch | LabeledBagOfTilesSampleBatch, ) -> None: - assert isinstance(trainer.logger, MLFlowLogger) sl_preds, tl_preds, batch_mask, batch_attention = outputs - _, metadata_batch = batch + metadata_batch = batch[-1] self.table["slide"].extend([m["slide_name"] for m in metadata_batch]) self.table["sl_prediction"].extend(sl_preds.tolist()) @@ -117,12 +118,40 @@ def on_predict_batch_end( artifact_path=str(mask_builder.save_dir), ) - def on_predict_epoch_end( - self, trainer: Trainer, pl_module: LightningModule + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, # type: ignore[override] + batch: LabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, ) -> None: + self._on_batch_end(trainer, outputs, batch) + + def on_predict_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, + batch: UnlabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + self._on_batch_end(trainer, outputs, batch) + + def _on_epoch_end(self) -> None: df = pd.DataFrame(self.table) df.to_json("sl_predictions.json", orient="split") mlflow.log_artifact( "sl_predictions.json", artifact_path="tables", ) + + def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + self._on_epoch_end() + + def on_predict_epoch_end( + self, trainer: Trainer, pl_module: LightningModule + ) -> None: + self._on_epoch_end() diff --git a/prostate_cancer/callbacks/nested_metrics_callback.py b/prostate_cancer/callbacks/nested_metrics_callback.py index c17cdf0b..23fea58b 100644 --- a/prostate_cancer/callbacks/nested_metrics_callback.py +++ b/prostate_cancer/callbacks/nested_metrics_callback.py @@ -1,45 +1,18 @@ from typing import Any -import lightning.pytorch as pl -import mlflow -import pandas as pd -from lightning import Callback -from rationai.mlkit.lightning.loggers import MLFlowLogger -from rationai.mlkit.metrics import NestedMetricCollection -from torchmetrics import ( - AUROC, - Accuracy, - NegativePredictiveValue, - Precision, - Recall, - Specificity, -) +from lightning import LightningModule, Trainer +from prostate_cancer.callbacks.nested_metrics_callback_base import ( + NestedMetricsCallbackBase, +) from prostate_cancer.typing import LabeledTileSampleBatch -class NestedMetricsCallback(Callback): - """Calculates metrics using the `NestedMetricCollection` in the test stage.""" - - def __init__(self, threshold: float) -> None: - # In the test mode, log metrics for each slide - self.nested_test_metrics = NestedMetricCollection( - metrics={ - "AUC": AUROC("binary"), - "accuracy": Accuracy("binary", threshold), - "precision": Precision("binary", threshold), - "recall": Recall("binary", threshold), - "specificity": Specificity("binary", threshold), - "negative_predictive_value": NegativePredictiveValue( - "binary", threshold - ), - } - ) - +class NestedMetricsCallback(NestedMetricsCallbackBase): def on_test_batch_end( self, - trainer: pl.Trainer, - pl_module: pl.LightningModule, + trainer: Trainer, + pl_module: LightningModule, outputs: Any, batch: LabeledTileSampleBatch, batch_idx: int, @@ -49,13 +22,3 @@ def on_test_batch_end( # Update slide-level metrics self.nested_test_metrics.update(outputs, targets, metadata["slide"]) - - def on_test_epoch_end( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - assert isinstance(trainer.logger, MLFlowLogger) - - metrics = self.nested_test_metrics.compute() - pd.DataFrame(metrics).to_json("nested_metrics.json", orient="split") - mlflow.log_artifact("nested_metrics.json") - self.nested_test_metrics.reset() diff --git a/prostate_cancer/callbacks/nested_metrics_callback_base.py b/prostate_cancer/callbacks/nested_metrics_callback_base.py new file mode 100644 index 00000000..635e646d --- /dev/null +++ b/prostate_cancer/callbacks/nested_metrics_callback_base.py @@ -0,0 +1,42 @@ +from abc import ABC + +import mlflow +import pandas as pd +from lightning import Callback, LightningModule, Trainer +from rationai.mlkit.lightning.loggers import MLFlowLogger +from rationai.mlkit.metrics import NestedMetricCollection +from torchmetrics import ( + AUROC, + Accuracy, + NegativePredictiveValue, + Precision, + Recall, + Specificity, +) + + +class NestedMetricsCallbackBase(Callback, ABC): + """Calculates metrics using the `NestedMetricCollection` in the test stage, grouped per slide.""" + + def __init__(self, threshold: float) -> None: + super().__init__() + self.nested_test_metrics = NestedMetricCollection( + metrics={ + "AUC": AUROC("binary"), + "accuracy": Accuracy("binary", threshold), + "precision": Precision("binary", threshold), + "recall": Recall("binary", threshold), + "specificity": Specificity("binary", threshold), + "negative_predictive_value": NegativePredictiveValue( + "binary", threshold + ), + } + ) + + def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + assert isinstance(trainer.logger, MLFlowLogger) + + metrics = self.nested_test_metrics.compute() + pd.DataFrame(metrics).to_json("nested_metrics.json", orient="split") + mlflow.log_artifact("nested_metrics.json") + self.nested_test_metrics.reset() diff --git a/prostate_cancer/callbacks/nested_metrics_callback_mil.py b/prostate_cancer/callbacks/nested_metrics_callback_mil.py new file mode 100644 index 00000000..954b1374 --- /dev/null +++ b/prostate_cancer/callbacks/nested_metrics_callback_mil.py @@ -0,0 +1,36 @@ +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.nested_metrics_callback_base import ( + NestedMetricsCallbackBase, +) +from prostate_cancer.typing import LabeledBagOfTilesSampleBatch, MILModelOutput + + +class NestedMetricsCallbackMIL(NestedMetricsCallbackBase): + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, # type: ignore[override] + batch: LabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + _, tl_targets, _, metadata = batch + _, tl_outputs_raw, mask, _ = outputs + + mask_bool = mask.bool() + tl_outputs_valid = tl_outputs_raw[mask_bool] + targets_valid = tl_targets[mask_bool] + + keys = [ + metadata[i]["slide_name"] + for i in range(mask_bool.shape[0]) + for j in range(mask_bool.shape[1]) + if mask_bool[i, j] + ] + + # Update slide-level metrics + self.nested_test_metrics.update( + tl_outputs_valid.cpu(), targets_valid.cpu(), keys + ) diff --git a/prostate_cancer/callbacks/slide_histograms_callback_mil.py b/prostate_cancer/callbacks/slide_histograms_callback_mil.py new file mode 100644 index 00000000..2b6e07e9 --- /dev/null +++ b/prostate_cancer/callbacks/slide_histograms_callback_mil.py @@ -0,0 +1,46 @@ +from typing import TYPE_CHECKING, cast + +import numpy as np +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.histograms_callback_base import ( + HistogramsCallbackBase, +) +from prostate_cancer.typing import MILModelOutput, UnlabeledBagOfTilesSampleBatch + + +if TYPE_CHECKING: + from prostate_cancer.datamodule import BagOfTilesDataModule + + +class SlideHistogramsCallbackMIL(HistogramsCallbackBase): + def setup( + self, trainer: Trainer, pl_module: LightningModule, stage: str | None = None + ) -> None: + if not hasattr(trainer, "datamodule"): + raise ValueError("Trainer should have datamodule attribute") + + datamodule = cast("BagOfTilesDataModule", trainer.datamodule) + slides = datamodule.predict.slides + self._slide_targets = dict(zip(slides["id"], slides["carcinoma"], strict=True)) + + def on_predict_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, + batch: UnlabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + sl_outputs, _, _, _ = outputs + _, metadata_batch = batch + + preds = sl_outputs.detach().cpu().numpy().flatten() + labels = np.array( + [float(self._slide_targets[m["slide_id"]]) for m in metadata_batch], + dtype=np.float32, + ) + + self.all_preds.append(preds) + self.all_labels.append(labels) diff --git a/prostate_cancer/callbacks/tile_histograms_callback_mil.py b/prostate_cancer/callbacks/tile_histograms_callback_mil.py new file mode 100644 index 00000000..0d409fe5 --- /dev/null +++ b/prostate_cancer/callbacks/tile_histograms_callback_mil.py @@ -0,0 +1,29 @@ +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.histograms_callback_base import ( + HistogramsCallbackBase, +) +from prostate_cancer.typing import LabeledBagOfTilesSampleBatch, MILModelOutput + + +class TileHistogramsCallbackMIL(HistogramsCallbackBase): + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: MILModelOutput, # type: ignore[override] + batch: LabeledBagOfTilesSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + _, y, _, _ = batch + _, tl_outputs_raw, mask, _ = outputs + + tl_outputs_valid = tl_outputs_raw[mask.bool()] + preds = tl_outputs_valid.detach().cpu().numpy().flatten() + + labels_valid = y[mask.bool()] + labels = labels_valid.detach().cpu().numpy().flatten() + + self.all_preds.append(preds) + self.all_labels.append(labels) diff --git a/prostate_cancer/callbacks/tile_histograms_callback_tl.py b/prostate_cancer/callbacks/tile_histograms_callback_tl.py new file mode 100644 index 00000000..ac028907 --- /dev/null +++ b/prostate_cancer/callbacks/tile_histograms_callback_tl.py @@ -0,0 +1,26 @@ +from typing import Any + +from lightning import LightningModule, Trainer + +from prostate_cancer.callbacks.histograms_callback_base import ( + HistogramsCallbackBase, +) +from prostate_cancer.typing import LabeledTileSampleBatch + + +class TileHistogramsCallbackTile(HistogramsCallbackBase): + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: LightningModule, + outputs: Any, + batch: LabeledTileSampleBatch, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + _, y, _ = batch + preds = outputs.detach().cpu().numpy().flatten() + labels = y.detach().cpu().numpy().flatten() + + self.all_preds.append(preds) + self.all_labels.append(labels)