[predict] Split predict into functions and implement in segmentation

src/mednet/libs/classification/scripts/experiment.py

@@ -111,11 +111,11 @@ def experiment(
 
     from .predict import predict
 
-    predictions_output = output_folder / "predictions.json"
+    predictions_output = output_folder / "predictions"
 
     ctx.invoke(
         predict,
-        output=predictions_output,
+        output_folder=predictions_output,
         model=model,
         datamodule=datamodule,
         device=device,
@@ -135,7 +135,7 @@ def experiment(
 
     ctx.invoke(
         evaluate,
-        predictions=predictions_output,
+        predictions=predictions_output / "predictions.json",
         output_folder=output_folder,
         threshold="validation",
     )

src/mednet/libs/classification/scripts/predict.py

@@ -34,7 +34,7 @@ logger = setup(__name__.split(".")[0], format="%(levelname)s: %(message)s")
 @reusable_options
 @verbosity_option(logger=logger, cls=ResourceOption, expose_value=False)
 def predict(
-    output,
+    output_folder,
     model,
     datamodule,
     batch_size,
@@ -56,22 +56,24 @@ def predict(
         setup_datamodule,
     )
 
+    predictions_file = output_folder / "predictions.json"
+    predictions_file.parent.mkdir(parents=True, exist_ok=True)
+
     setup_datamodule(datamodule, model, batch_size, parallel)
     model = load_checkpoint(model, weight)
     device_manager = DeviceManager(device)
-    save_json_data(datamodule, model, output, device_manager)
+    save_json_data(datamodule, model, predictions_file, device_manager)
 
     predictions = run(model, datamodule, device_manager)
 
-    output.parent.mkdir(parents=True, exist_ok=True)
-    if output.exists():
-        backup = output.parent / (output.name + "~")
+    if predictions_file.exists():
+        backup = predictions_file.parent / (predictions_file.name + "~")
         logger.warning(
-            f"Output predictions file `{str(output)}` exists - "
+            f"Output predictions file `{str(predictions_file)}` exists - "
             f"backing it up to `{str(backup)}`...",
         )
-        shutil.copy(output, backup)
+        shutil.copy(predictions_file, backup)
 
-    with output.open("w") as f:
+    with predictions_file.open("w") as f:
         json.dump(predictions, f, indent=2)
-    logger.info(f"Predictions saved to `{str(output)}`")
+    logger.info(f"Predictions saved to `{str(predictions_file)}`")

src/mednet/libs/classification/tests/test_cli_classification.py

@@ -329,7 +329,7 @@ def test_predict_pasa_montgomery(temporary_basedir, datadir):
     runner = CliRunner()
 
     with stdout_logging() as buf:
-        output = temporary_basedir / "predictions.json"
+        output = temporary_basedir / "predictions"
         last = _get_checkpoint_from_alias(
             temporary_basedir / "results",
             "periodic",
@@ -343,7 +343,7 @@ def test_predict_pasa_montgomery(temporary_basedir, datadir):
                 "-vv",
                 "--batch-size=1",
                 f"--weight={str(last)}",
-                f"--output={str(output)}",
+                f"--output-folder={str(output)}",
             ],
         )
         _assert_exit_0(result)
@@ -379,7 +379,8 @@ def test_evaluate_pasa_montgomery(temporary_basedir):
     runner = CliRunner()
 
     with stdout_logging() as buf:
-        prediction_path = temporary_basedir / "predictions.json"
+        prediction_path = temporary_basedir / "predictions"
+        predictions_file = prediction_path / "predictions.json"
         evaluation_filename = "evaluation.json"
         evaluation_file = temporary_basedir / evaluation_filename
         result = runner.invoke(
@@ -387,7 +388,7 @@ def test_evaluate_pasa_montgomery(temporary_basedir):
             [
                 "-vv",
                 "montgomery",
-                f"--predictions={str(prediction_path)}",
+                f"--predictions={predictions_file}",
                 f"--output-folder={str(temporary_basedir)}",
                 "--threshold=test",
             ],
@@ -440,9 +441,11 @@ def test_experiment(temporary_basedir):
     _assert_exit_0(result)
 
     assert (output_folder / "model" / "meta.json").exists()
-    assert (output_folder / "model" / f"model-at-epoch={num_epochs-1}.ckpt").exists()
-    assert (output_folder / "predictions.json").exists()
-    assert (output_folder / "predictions.meta.json").exists()
+    assert (
+        output_folder / "model" / f"model-at-epoch={num_epochs-1}.ckpt"
+    ).exists()
+    assert (output_folder / "predictions" / "predictions.json").exists()
+    assert (output_folder / "predictions" / "predictions.meta.json").exists()
 
     # Need to glob because we cannot be sure of the checkpoint with lowest validation loss
     assert (

src/mednet/libs/common/scripts/predict.py

@@ -32,20 +32,18 @@ def reusable_options(f):
     """
 
     @click.option(
-        "--output",
+        "--output-folder",
         "-o",
-        help="""Path to a JSON file in which to save predictions for all samples in the
-        input DataModule (leading directories are created if they do not
-        exist).""",
+        help="Directory in which to save predictions (created if does not exist)",
         required=True,
-        default="predictions.json",
-        cls=ResourceOption,
         type=click.Path(
-            file_okay=True,
-            dir_okay=False,
+            file_okay=False,
+            dir_okay=True,
             writable=True,
             path_type=pathlib.Path,
         ),
+        default="predictions",
+        cls=ResourceOption,
     )
     @click.option(
         "--model",

src/mednet/libs/segmentation/models/separate.py

@@ -5,7 +5,6 @@
 
 import typing
 
-import torch
 from mednet.libs.common.data.typing import Sample
 
 from .typing import BinaryPrediction, MultiClassPrediction
@@ -27,7 +26,7 @@ def _as_predictions(
     list[BinaryPrediction | MultiClassPrediction]
         A list of typed predictions that can be saved to disk.
     """
-    return [(v[1]["name"], v[1]["target"], v[0].item()) for v in samples]
+    return [(v[1]["name"], v[1]["target"], v[0]) for v in samples]
 
 
 def separate(batch: Sample) -> list[BinaryPrediction | MultiClassPrediction]:
@@ -58,4 +57,5 @@ def separate(batch: Sample) -> list[BinaryPrediction | MultiClassPrediction]:
         {key: value[i] for key, value in batch[1].items()}
         for i in range(len(batch[0]))
     ]
-    return _as_predictions(zip(torch.flatten(batch[0]), metadata))
+
+    return _as_predictions(zip(batch[0], metadata))

src/mednet/libs/segmentation/scripts/experiment.py

@@ -112,11 +112,11 @@ def experiment(
 
     from .predict import predict
 
-    predictions_output = output_folder / "predictions.json"
+    predictions_output = output_folder / "predictions"
 
     ctx.invoke(
         predict,
-        output=predictions_output,
+        output_folder=predictions_output,
         model=model,
         datamodule=datamodule,
         device=device,
@@ -131,7 +131,7 @@ def experiment(
         f"Prediction runtime: {predict_stop_timestamp-predict_start_timestamp}"
     )
 
-    """evaluation_start_timestamp = datetime.now()
+    evaluation_start_timestamp = datetime.now()
     logger.info(f"Started evaluation at {evaluation_start_timestamp}")
 
     from .evaluate import evaluate
@@ -140,14 +140,15 @@ def experiment(
         evaluate,
         predictions=predictions_output,
         output_folder=output_folder,
-        threshold="validation",
+        # threshold="validation",
+        threshold=0.5,
     )
 
     evaluation_stop_timestamp = datetime.now()
     logger.info(f"Ended prediction in {evaluation_stop_timestamp}")
     logger.info(
         f"Prediction runtime: {evaluation_stop_timestamp-evaluation_start_timestamp}"
-    )"""
+    )
 
     experiment_stop_timestamp = datetime.now()
     logger.info(

src/mednet/libs/segmentation/scripts/predict.py

@@ -3,15 +3,47 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 
+import pathlib
+
 import click
+import h5py
+import PIL
 from clapper.click import ResourceOption, verbosity_option
 from clapper.logging import setup
 from mednet.libs.common.scripts.click import ConfigCommand
 from mednet.libs.common.scripts.predict import reusable_options
+from tqdm import tqdm
 
 logger = setup(__name__.split(".")[0], format="%(levelname)s: %(message)s")
 
 
+def _save_hdf5(
+    stem: pathlib.Path, prob: PIL.Image.Image, output_folder: pathlib.Path
+):
+    """Save prediction maps as image in the same format as the test image.
+
+    Parameters
+    ----------
+    stem
+        Name of the file without extension on the original dataset.
+
+    prob
+        Monochrome Image with prediction maps.
+
+    output_folder
+        Directory in which to store predictions.
+    """
+
+    fullpath = output_folder / f"{stem}.hdf5"
+    tqdm.write(f"Saving {fullpath}...")
+    fullpath.parent.mkdir(parents=True, exist_ok=True)
+    with h5py.File(fullpath, "w") as f:
+        data = prob.squeeze(0).numpy()
+        f.create_dataset(
+            "array", data=data, compression="gzip", compression_opts=9
+        )
+
+
 @click.command(
     entry_point_group="mednet.libs.segmentation.config",
     cls=ConfigCommand,
@@ -34,7 +66,7 @@ logger = setup(__name__.split(".")[0], format="%(levelname)s: %(message)s")
 @reusable_options
 @verbosity_option(logger=logger, cls=ResourceOption, expose_value=False)
 def predict(
-    output,
+    output_folder,
     model,
     datamodule,
     batch_size,
@@ -45,9 +77,6 @@ def predict(
 ) -> None:  # numpydoc ignore=PR01
     """Run inference (generates scores) on all input images, using a pre-trained model."""
 
-    import json
-    import shutil
-
     from mednet.libs.classification.engine.predictor import run
     from mednet.libs.common.engine.device import DeviceManager
     from mednet.libs.common.scripts.predict import (
@@ -56,31 +85,15 @@ def predict(
         setup_datamodule,
     )
 
+    predictions_file = output_folder / "predictions.json"
+
     setup_datamodule(datamodule, model, batch_size, parallel)
     model = load_checkpoint(model, weight)
     device_manager = DeviceManager(device)
-    save_json_data(datamodule, model, output, device_manager)
+    save_json_data(datamodule, model, predictions_file, device_manager)
 
     predictions = run(model, datamodule, device_manager)
 
-    output.parent.mkdir(parents=True, exist_ok=True)
-    if output.exists():
-        backup = output.parent / (output.name + "~")
-        logger.warning(
-            f"Output predictions file `{str(output)}` exists - "
-            f"backing it up to `{str(backup)}`...",
-        )
-        shutil.copy(output, backup)
-
-    # Remove targets from predictions, as they are images and not serializable
-    # A better solution should be found
-    serializable_predictions = {}
-    for split_name, sample in predictions.items():
-        split_predictions = []
-        for s in sample:
-            split_predictions.append((s[0], s[2]))
-        serializable_predictions[split_name] = split_predictions
-
-    with output.open("w") as f:
-        json.dump(serializable_predictions, f, indent=2)
-    logger.info(f"Predictions saved to `{str(output)}`")
+    for split_name, split in predictions.items():
+        for sample in split:
+            _save_hdf5(sample[0], sample[2], output_folder)