cls_eval.py

import os
from functools import partial

import pandas as pd
import torch
import torch.nn.functional as F
from monai.transforms import SaveImaged

from engine import DefaultEvaluator, parse_args_and_launch
from metrics import ConfMatMetric, AUROCMetric
from prob_aggregator import ProbabilityAggregator
from utils import sliding_patch_inference, get_filename_from_meta, name_with_key_formatter, load_config


class ClsEvaluator(DefaultEvaluator):
    def __init__(self, cfg, **kwargs):
        super().__init__(cfg, **kwargs)
        self.prob_aggregator = ProbabilityAggregator(cfg.eval.prob_aggregator)
        self.category_names = self.eval_loader.category_names
        self.num_classes = len(self.category_names["CLS"])

    @classmethod
    def _build_metrics(cls, cfg, category_names):
        return [
            # Using patch-level f1-score macro average as main metric (for model selection)
            ConfMatMetric(name="Patch_ConfMat", category_names=category_names["CLS"],
                          metric_name=["f1_score", "precision", "recall", "accuracy"]),
            AUROCMetric(name="Patch_AUROC", category_names=category_names["CLS"]),
            ConfMatMetric(name="ConfMat", category_names=category_names["CLS"],
                          metric_name=["f1_score", "precision", "recall", "accuracy"]),
            AUROCMetric(name="AUROC", category_names=category_names["CLS"]),
        ]

    def _get_patch_label(self, inputs, labels):
        # Assuming inputs are binary masks of shape: (B*b, 1, h, w, d) and labels are global labels of shape: (B)
        # where b = sliding window batch size, and h, w, d = patch dimensions
        input_volume = inputs[0].numel()
        all_except_batch = tuple(range(1, inputs.dim()))
        p_labels = torch.gt(
            inputs.sum(dim=all_except_batch) / input_volume, self.cfg.eval.sw_patch_fg_threshold)  # Shape: (B*b)
        p_labels = p_labels.view(labels.size(0), -1) * labels  # Shape: (B, b)
        return p_labels.view(-1)  # Shape: (B*b)

    def _save_batch_results(self, batch_data, patch_labels, patch_probs, volume_probs):

        output_dir = os.path.join(self.cfg.run_dir, "results")
        csv_path = os.path.join(output_dir, 'volume_preds.csv')

        # Convert to regular Tensors (no embedded meta), otherwise saver will ignore provided metadata.
        patch_labels = torch.Tensor(patch_labels)
        patch_probs = torch.Tensor(patch_probs)

        # Dynamically create separate keys for each label and probability channel
        saver_keys = [f"label_{c}" for c in range(self.num_classes)] + \
            [f"prob_{c}" for c in range(self.num_classes)]

        saver = SaveImaged(
            keys=saver_keys,
            meta_keys=None,
            meta_key_postfix="meta_dict",
            output_postfix="",
            output_dir=output_dir,
            output_ext=".nii.gz",
            output_name_formatter=name_with_key_formatter,
            data_root_dir=self.cfg.dataset.data_dir,
            separate_folder=False,
            print_log=False,
            mode='nearest'
        )

        batch_size = patch_labels.shape[0]
        rows = []

        for i in range(batch_size):
            meta_dict = dict(batch_data['image'][i].meta)
            sample_dict = {}
            for c, cat_name in enumerate(self.category_names["CLS"]):
                # Create a separate meta for each channel
                label_meta = {**meta_dict, "key": f"{cat_name}_label"}
                prob_meta = {**meta_dict, "key": f"{cat_name}_prob"}
                sample_dict[f"label_{c}"] = patch_labels[i][..., c]
                sample_dict[f"label_{c}_meta_dict"] = label_meta
                sample_dict[f"prob_{c}"] = patch_probs[i][..., c]
                sample_dict[f"prob_{c}_meta_dict"] = prob_meta
            saver(sample_dict)
            filename = os.path.basename(get_filename_from_meta(meta_dict))
            rows.append([filename] + volume_probs[i].tolist())

        columns = ["File Name"] + self.category_names["CLS"]
        pd.DataFrame(rows, columns=columns).to_csv(
            csv_path, index=False, mode='a', header=not os.path.exists(csv_path))

    @torch.no_grad()
    def _evaluate_batch(self, batch_data):
        self.model.eval()
        images = batch_data['image'].to(self.device)  # Shape: (B, C, H, W, D)
        labels = batch_data['label'].to(self.device)  # Shape: (B)
        masks = batch_data['mask'].to(self.device)  # Shape: (B, 1, H, W, D)

        with torch.amp.autocast(self.device.type):
            # Perform sliding window inference to find the label of each patch
            patch_labels = sliding_patch_inference(
                inputs=masks,
                roi_size=self.cfg.eval.sw_patch_size,
                sw_batch_size=self.cfg.eval.sw_batch_size,
                predictor=partial(self._get_patch_label, labels=labels),
                overlap=self.cfg.eval.sw_overlap
            ).squeeze(-1)  # Shape: (B, H_p, W_p, D_p), where H_p, W_p, D_p = number of patches in each dimension
            # Perform sliding window inference to find model prediction for each patch
            patch_logits = sliding_patch_inference(
                inputs=images,
                roi_size=self.cfg.eval.sw_patch_size,
                sw_batch_size=self.cfg.eval.sw_batch_size,
                predictor=self.model,
                overlap=self.cfg.eval.sw_overlap
            )  # Shape: (B, H_p, W_p, D_p, N), where N = number of classes

            # Convert logits to probabilities using softmax
            patch_probs = F.softmax(patch_logits, dim=-1)  # Shape: (B, H_p, W_p, D_p, N)
            if self.cfg.eval.binary_category and patch_probs.size(1) > 2:
                # Keep normal probability (index 0) and sum all abnormal probabilities (index 1+)
                abnormal_probs = torch.sum(patch_probs[..., 1:], dim=-1, keepdim=True)  # Sum abnormal classes
                normal_probs = patch_probs[..., 0:1]  # Keep normal class probability
                # Combine into binary probabilities [normal, abnormal]
                patch_probs = torch.cat([normal_probs, abnormal_probs], dim=-1)  # Shape: (B, H_p, W_p, D_p, 2)

            # Get patch class predictions and patch_labels in one-hot format
            patch_preds = torch.argmax(patch_probs, dim=-1)  # Shape: (B, H_p, W_p, D_p)
            patch_preds = F.one_hot(patch_preds, num_classes=self.num_classes)  # Shape: (B, H_p, W_p, D_p, N)
            patch_labels = F.one_hot(patch_labels, num_classes=self.num_classes)  # Shape: (B, H_p, W_p, D_p, N)

            volume_probs = self.prob_aggregator.get_aggregated_probs(patch_probs)  # Shape: (B, N)
            volume_preds = torch.argmax(volume_probs, dim=1)  # Shape: (B)
            volume_preds = F.one_hot(volume_preds, num_classes=self.num_classes)  # Shape: (B, N)
            labels = F.one_hot(labels, num_classes=self.num_classes)  # Shape: (B, N)

            for metric in self.metrics:
                if metric.name.startswith("Patch"):
                    # Compute patch-level metrics
                    if metric.require_binarized_pred:
                        metric.accumulate(patch_preds.view(-1, self.num_classes),
                                          patch_labels.view(-1, self.num_classes))
                    else:
                        metric.accumulate(patch_probs.view(-1, self.num_classes),
                                          patch_labels.view(-1, self.num_classes))
                else:
                    # Compute volume-level metrics
                    if metric.require_binarized_pred:
                        metric.accumulate(volume_preds, labels)
                    else:
                        metric.accumulate(volume_probs, labels)

        if self.output_results:
            self._save_batch_results(batch_data, patch_labels, patch_probs, volume_probs)


def main(args):
    cfg = load_config(args.config_file, args.opts)
    evaluator = ClsEvaluator(
        cfg, debug=args.debug, output_results=True, ckpt_path=args.ckpt_path, no_progress=args.no_progress)
    evaluator.logger.log_params(cfg, args)
    evaluator.evaluate()


if __name__ == "__main__":
    parse_args_and_launch(main)  # pragma: no cover