Inference with a safe model artifact

#all_skip
from fastMONAI.vision_all import *
from monai.apps import DecathlonDataset

tasks = {
    "brain": "Task01_BrainTumour",
    "heart": "Task02_Heart", 
    "spleen": "Task09_Spleen"
}

task = tasks["heart"]

model_artifact_path = Path(f"model_artifacts/{task}")

path = Path('../data')
path.mkdir(exist_ok=True)

test_data = DecathlonDataset(root_dir=path, task=task, section="test", download=True,
                                 cache_num=0, num_workers=3)

2025-08-29 14:42:58,879 - INFO - Verified 'Task02_Heart.tar', md5: 06ee59366e1e5124267b774dbd654057.
2025-08-29 14:42:58,879 - INFO - File exists: ../data/Task02_Heart.tar, skipped downloading.
2025-08-29 14:42:58,880 - INFO - Non-empty folder exists in ../data/Task02_Heart, skipped extracting.

test_imgs = [data['image'] for data in test_data.data]

import mlflow

# Leave as None to use the newest completed tutorial run.
MODEL_RUN_ID = None
# MODEL_RUN_ID = "your-mlflow-run-id"

if MODEL_RUN_ID is None:
    experiment = mlflow.get_experiment_by_name(task)
    if experiment is None:
        raise RuntimeError(f"No MLflow experiment found for {task!r}. Run tutorial 11d first.")
    runs = mlflow.search_runs(
        experiment_ids=[experiment.experiment_id],
        filter_string="attributes.status = 'FINISHED'",
        order_by=["attributes.start_time DESC"],
        max_results=1,
    )
    if runs.empty:
        raise RuntimeError(f"No completed MLflow run found for {task!r}. Run tutorial 11d first.")
    MODEL_RUN_ID = runs.iloc[0]["run_id"]
    print(f"Using latest completed run: {MODEL_RUN_ID}")

model_path = find_model_artifacts(
    run_ids={"model": MODEL_RUN_ID},
    artifact_role="best",
)["model"]

import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_inf, inference_config = load_model_resources(model_path, device=device)

(
    inference_config['apply_reorder'],
    inference_config['target_spacing'],
    inference_config['output'],
)

save_path = Path(f'../data/{task}/pred_masks')
save_path.mkdir(parents=True, exist_ok=True)

idx = 3
img_fn = test_imgs[idx]
img_fn

'../data/Task02_Heart/imagesTs/la_001.nii.gz'

pred_fn = inference_from_config(
    model_inf,
    inference_config,
    fn=img_fn,
    save_path=save_path,
)
pred_fn

from torchio import Subject, ScalarImage, LabelMap

subject = Subject(image=ScalarImage(img_fn), mask=LabelMap(pred_fn))
subject.plot(figsize=(10,5))