results_computing.py

# Copyright (c) IRCAD France
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree. (GNU GPL v3)
import argparse
import pickle
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import SimpleITK as sitk
import torch

from SLIP_model.inference_module import Inference_module


def calculate_dice_score(pred, gt):
    """
    Calculate Dice score between prediction and ground truth.
    Args:
        pred: Predicted mask tensor
        gt: Ground truth mask tensor
    Returns:
        dice_score: Float value between 0 and 1
    """
    pred = (pred.float() > 0).float()
    gt = (gt.float() > 0).float()

    intersection = (pred * gt).sum()
    union = pred.sum() + gt.sum()

    if union == 0:
        return 1.0 if intersection == 0 else 0.0

    dice = (2.0 * intersection) / union
    return dice.item()


def region_points_3d(pred_masks, gt_masks, num_pt=1):
    """
    Sample points from the largest 3D error region (FP or FN).

    Args:
        pred_masks: Predicted masks tensor with shape [B, C, Z, H, W]
        gt_masks: Ground truth masks tensor with shape [B, C, Z, H, W]
        num_pt: Number of points to sample per batch

    Returns:
        points: Sampled points coordinates [B, num_pt, 3] (z, y, x)
        labels: Labels for sampled points [B, num_pt] (0 for FP, 1 for FN)
    """
    gt_masks = gt_masks > 0
    pred_masks = pred_masks > 0

    B, C, Z, H, W = gt_masks.shape
    device = gt_masks.device

    points = torch.zeros((B, num_pt, 3), dtype=torch.float, device=device)
    labels = torch.full((B, num_pt), -1, dtype=torch.int32, device=device)

    for b in range(B):
        # Get 3D masks for this batch (collapse channels if needed)
        gt_3d = gt_masks[b]
        pred_3d = pred_masks[b]

        if C > 1:
            gt_3d = torch.any(gt_3d, dim=0)  # [Z, H, W]
            pred_3d = torch.any(pred_3d, dim=0)  # [Z, H, W]
        else:
            gt_3d = gt_3d.squeeze(0)
            pred_3d = pred_3d.squeeze(0)

        # Calculate 3D error regions
        fp_mask = ~gt_3d & pred_3d  # False Positives
        fn_mask = gt_3d & ~pred_3d  # False Negatives

        # Choose the larger area
        fp_area = fp_mask.sum().item()
        fn_area = fn_mask.sum().item()

        selected_mask = fp_mask if fp_area >= fn_area else fn_mask
        selected_label = 0 if fp_area >= fn_area else 1

        # Get 3D coordinates of error region
        z_coords, y_coords, x_coords = torch.where(selected_mask)

        for i in range(num_pt):
            if len(z_coords) > 0:
                idx = torch.randint(0, len(z_coords), (1, ), device=device).item()

                points[b, i, 0] = z_coords[idx].item()
                points[b, i, 1] = y_coords[idx].item()
                points[b, i, 2] = x_coords[idx].item()
                labels[b, i] = selected_label

    return points, labels


def process_image_with_clicks(image_path, gt_path, click_number=4, target_label=None, precomputed_gt_binary=None):
    """
    Process a single image with interactive clicks and track Dice scores and timing.

    Click selection strategy: first click targets the largest error region (FN or FP)
    relative to the current prediction, updated after every click.

    Args:
        image_path: Path to input NIfTI image.
        gt_path: Path to ground truth NIfTI label map (may be None if precomputed_gt_binary is given).
        click_number: Total number of simulated clicks to perform.
        target_label: Integer label to segment, or None to auto-select the first non-zero label.
        precomputed_gt_binary: Optional pre-binarised GT numpy array (float32, shape Z×H×W)
                               that bypasses NIfTI loading for 4D inputs.

    Returns:
        dice_scores: List of Dice scores after each click.
        click_times: List of wall-clock times (seconds) per click.
        final_prediction: Final segmentation tensor in original image space.
    """

    # ── Let the model load and preprocess the image ───────────────────────
    predict.process_image(
        image_path=image_path,
        gt_path=gt_path,
        target_label=target_label,
        precomputed_gt_binary=precomputed_gt_binary,
    )

    # Derive shape and GT tensor directly from the model's loaded data
    # gt3D_original: (1, 1, Z, H, W) in original image space
    gt_tensor = predict.gt3D_original.squeeze()  # (Z, H, W)
    Z, H, W = gt_tensor.shape

    dice_scores = []
    click_times = []

    print(f"\nProcessing {image_path.name}")
    print(f"Volume shape (original): {predict.vol_original.shape}, GT shape: {gt_tensor.shape}")

    # Subsequent frames will be selected based on error after each slice is processed
    total_clicks = click_number
    click_count = 0
    pred_masks = torch.zeros(Z, H, W, dtype=torch.float32, device=predict.device)
    try:
        # Perform clicks on the selected slice
        for _ in range(total_clicks):
            click_count += 1

            # Prepare tensors for region_points function
            pred_tensor = pred_masks.unsqueeze(0).unsqueeze(0).float()  # [1, 1, Z, H, W]
            gt_tensor_batch = gt_tensor.unsqueeze(0).unsqueeze(0).float()  # [1, 1, Z, H, W]

            # Sample point
            points, labels = region_points_3d(pred_tensor, gt_tensor_batch)

            # Extract point coordinates (z, y, x)
            point_coords = (
                int(points[0, 0, 0].item()),  # z
                int(points[0, 0, 1].item()),  # y
                int(points[0, 0, 2].item())  # x
            )
            label = labels[0, 0].item()
            include = (label == 1)  # True for FN (positive), False for FP (negative)
            all_clicks = [[int(points[0, i, 0].item()),
                           int(points[0, i, 1].item()),
                           int(points[0, i, 2].item())] for i in range(points.shape[1])]
            all_labels = [int(labels[0, i].item()) for i in range(labels.shape[1])]

            # Start timing
            start_time = time.time()
            pred_masks = predict.click_inference(all_clicks, all_labels)

            # End timing
            end_time = time.time()
            click_time = end_time - start_time
            click_times.append(click_time)

            # Calculate Dice score
            dice = calculate_dice_score(pred_masks, gt_tensor)
            dice_scores.append(dice)

            print(f"Click {click_count}/{total_clicks} -, "
                  f"Point: {point_coords}, Label: {'Positive' if include else 'Negative'}, "
                  f"Dice: {dice:.4f}, Time: {click_time:.3f}s")
            pred_masks = pred_masks.squeeze().float()
        # Get final prediction
        final_prediction = pred_masks
        predict.reset()
        return dice_scores, click_times, final_prediction
    finally:
        # ✅ Clean up all GPU tensors held by the model between images
        del pred_masks
        if hasattr(predict, 'image_embedding'):
            del predict.image_embedding
        if hasattr(predict, 'vol'):
            del predict.vol
        if hasattr(predict, 'gt3D'):
            del predict.gt3D
        torch.cuda.empty_cache()


def plot_mean_dice_progression(all_results, output_folder):
    """
    Plot mean Dice score across all images for each click.
    Args:
        all_results: Dictionary containing results for all images
        output_folder: Path to save the plot
    """
    output_folder = Path(output_folder)

    # Collect all dice scores per click position
    max_clicks = max(len(results['dice_scores']) for results in all_results.values())
    dice_per_click = [[] for _ in range(max_clicks)]

    for img_name, results in all_results.items():
        for click_idx, dice in enumerate(results['dice_scores']):
            dice_per_click[click_idx].append(dice)

    # Calculate mean and std for each click
    mean_dice = [np.mean(scores) if scores else 0 for scores in dice_per_click]
    std_dice = [np.std(scores) if scores else 0 for scores in dice_per_click]
    click_numbers = list(range(1, len(mean_dice) + 1))

    # Create the plot
    plt.figure(figsize=(12, 7))

    # Plot mean line with confidence interval
    plt.plot(click_numbers, mean_dice, 'b-o', linewidth=2, markersize=6, label='Mean Dice')
    plt.fill_between(click_numbers,
                     np.array(mean_dice) - np.array(std_dice),
                     np.array(mean_dice) + np.array(std_dice),
                     alpha=0.3,
                     color='blue',
                     label='±1 Std Dev')

    # Add individual image trajectories (lighter lines)
    for img_name, results in all_results.items():
        plt.plot(range(1,
                       len(results['dice_scores']) + 1),
                 results['dice_scores'],
                 alpha=0.2,
                 linewidth=1,
                 color='gray')

    # Formatting
    plt.xlabel('Click Number', fontsize=12, fontweight='bold')
    plt.ylabel('Dice Score', fontsize=12, fontweight='bold')
    plt.title('Mean Dice Score Progression Across All Images', fontsize=14, fontweight='bold')
    plt.grid(True, alpha=0.3, linestyle='--')
    plt.legend(fontsize=10)
    plt.ylim([0, 1.05])
    plt.xlim([0.5, len(mean_dice) + 0.5])

    # Add annotations for first and last mean dice
    if mean_dice:
        plt.annotate(f'Start: {mean_dice[0]:.3f}',
                     xy=(1, mean_dice[0]),
                     xytext=(1, mean_dice[0] - 0.1),
                     fontsize=10,
                     fontweight='bold',
                     bbox=dict(boxstyle='round,pad=0.5', facecolor='yellow', alpha=0.7))

        plt.annotate(f'End: {mean_dice[-1]:.3f}',
                     xy=(len(mean_dice), mean_dice[-1]),
                     xytext=(len(mean_dice), mean_dice[-1] + 0.05),
                     fontsize=10,
                     fontweight='bold',
                     bbox=dict(boxstyle='round,pad=0.5', facecolor='green', alpha=0.7))

    plt.tight_layout()

    # Save the plot
    plot_path = output_folder / "mean_dice_progression.png"
    plt.savefig(plot_path, dpi=300, bbox_inches='tight')
    print(f"Plot saved to {plot_path}")

    # Also save as PDF for publication quality
    plot_path_pdf = output_folder / "mean_dice_progression.pdf"
    plt.savefig(plot_path_pdf, bbox_inches='tight')
    print(f"Plot saved to {plot_path_pdf}")

    plt.close()

    # Print statistics
    print("\n" + "=" * 60)
    print("DICE PROGRESSION STATISTICS")
    print("=" * 60)
    for i, (mean, std) in enumerate(zip(mean_dice, std_dice), 1):
        print(f"Click {i:2d}: Mean Dice = {mean:.4f} ± {std:.4f}")

    if mean_dice:
        improvement = mean_dice[-1] - mean_dice[0]
        print(f"\nOverall Improvement: {improvement:.4f} ({improvement/mean_dice[0]*100:.1f}%)")


def load_gt_channels(gt_path: Path):
    """
    Returns (channels_array, label_ids, is_4d)

    For 4D/vector GT:
        channels_array : np.ndarray  (N, Z, Y, X)  — one binary mask per segment
        label_ids      : dict {int channel_idx → str label_name}
        is_4d          : True

    For 3D scalar GT:
        channels_array : np.ndarray  (Z, Y, X)
        label_ids      : list[int]   unique non-zero label values
        is_4d          : False
    """
    gt_sitk = sitk.ReadImage(str(gt_path))
    gt_arr = sitk.GetArrayFromImage(gt_sitk)
    ndim = gt_sitk.GetDimension()

    print(f"  GT raw array shape : {gt_arr.shape}", flush=True)
    print(f"  GT sitk dimension  : {ndim}", flush=True)

    # ── Case A : genuine 4D scalar NIfTI → (N, Z, Y, X) ──────────────────
    if ndim == 4:
        n_channels = gt_arr.shape[0]
        label_ids = {i: str(i) for i in range(n_channels)}
        return gt_arr, label_ids, True

    # ── Case B : 3D vector NIfTI → (Z, Y, X, N_groups) ───────────────────
    if ndim == 3 and gt_arr.ndim == 4:
        channels = []  # list of (Z, Y, X) binary arrays, indexed 0..N-1
        label_ids = {}  # {channel_idx (int) → label_name (str)}
        ch_idx = 0

        n_groups = gt_arr.shape[-1]
        for g in range(n_groups):
            group_vol = gt_arr[..., g]  # (Z, Y, X)
            unique_vals = np.unique(group_vol)
            unique_vals = unique_vals[unique_vals != 0]

            for val in sorted(unique_vals):
                channels.append((group_vol == val).astype(np.uint8))
                label_ids[ch_idx] = f"g{g}_v{int(val)}"
                ch_idx += 1

        channels_arr = np.stack(channels, axis=0)  # (N, Z, Y, X)
        print(f"  Expanded {n_groups} groups → {ch_idx} segments")
        print(f"  Segment IDs: {label_ids}")
        return channels_arr, label_ids, True

    # ── Case C : classic 3D integer label map → (Z, Y, X) ────────────────
    unique = np.unique(gt_arr)
    non_zero = unique[unique != 0].tolist()
    return gt_arr, non_zero, False


def process_folder(image_folder, gt_folder, output_folder, click_number, labels_to_process=None):
    """
    Process all images in a folder.
    Args:
        image_folder: Path to folder containing input images
        gt_folder: Path to folder containing ground truth masks
        output_folder: Path to save results
        num_slices: Number of slices to sample clicks from
        clicks_per_slice: Number of clicks per slice
        labels_to_process: List of specific labels to process (None = all labels)
    """
    image_folder = Path(image_folder)
    gt_folder = Path(gt_folder)
    output_folder = Path(output_folder)
    output_folder.mkdir(parents=True, exist_ok=True)

    # Find all image files
    image_files = sorted(list(image_folder.glob("*.nii.gz")) + list(image_folder.glob("*.nii")))

    if len(image_files) == 0:
        print(f"No images found in {image_folder}")
        return

    all_results = {}
    all_click_times = []
    all_clicks_data = []  # NEW: Store individual click data
    label_results = {}

    for img_path in image_files:
        # Find corresponding ground truth
        gt_path = gt_folder / img_path.name
        if not gt_path.exists():
            # Try alternative naming conventions
            gt_path = gt_folder / img_path.name.replace("_image", "_label")
            if not gt_path.exists():
                raise ValueError(f"Ground truth not found for image '{img_path.name}'. "
                                 f"Looked for '{img_path.name}' and "
                                 f"'{img_path.name.replace('_image', '_label')}' in '{gt_folder}'.")
        # ── Load GT and resolve label list ────────────────────────────────
        gt_arr, available_ids, is_4d = load_gt_channels(gt_path)

        if is_4d:
            # label_ids is now {int → str}, iterate over integer keys
            all_channel_ids = list(available_ids.keys())

            # gt_arr[ch] works because ch is always an int now
            non_empty_ids = [ch for ch in all_channel_ids if gt_arr[ch].any()]

            mode_str = (f"4D NIfTI — {gt_arr.shape[0]} channels, "
                        f"{len(non_empty_ids)} non-empty, "
                        f"names: {[available_ids[ch] for ch in non_empty_ids]}")
        else:
            non_empty_ids = available_ids
            mode_str = f"3D NIfTI — labels {non_empty_ids}"
        # Apply labels_to_process filter
        if labels_to_process is not None:
            target_ids = [l for l in labels_to_process if l in non_empty_ids]
        else:
            target_ids = non_empty_ids

        print(f"\n{'='*80}")
        print(f"Processing : {img_path.name}")
        print(f"GT mode    : {mode_str}", flush=True)
        print(f"Target IDs : {target_ids}")
        print(f"{'='*80}")

        # Process each label separately
        for label_idx in target_ids:
            print(f"\n--- Processing Label {label_idx} ---")
            if is_4d:
                binary_mask = gt_arr[label_idx].astype(np.float32)  # (Z, Y, X)
                precomputed = binary_mask
            else:
                precomputed = None  # process_image will do extraction itself

            try:
                dice_scores, click_times, final_pred = process_image_with_clicks(
                    img_path,
                    gt_path,
                    click_number,
                    target_label=label_idx,
                    precomputed_gt_binary=precomputed,
                )

                if len(dice_scores) == 0:
                    print(f"No valid clicks for {img_path.name} label {label_idx}, skipping...")
                    continue

                total_time = np.sum(click_times)
                result_key = f"{img_path.name}_label_{label_idx}"

                all_results[result_key] = {
                    'dice_scores': dice_scores,
                    'click_times': click_times,
                    'final_dice': dice_scores[-1],
                    'mean_click_time': np.mean(click_times),
                    'total_time': total_time,
                    'label': label_idx,
                    'image_name': img_path.name
                }

                # NEW: Store individual click data
                for click_idx, (dice, t) in enumerate(zip(dice_scores, click_times)):
                    all_clicks_data.append({
                        'image_name': img_path.name,
                        'label': label_idx,
                        'click_number': click_idx + 1,
                        'dice_score': dice,
                        'click_time': t,
                        'cumulative_time': np.sum(click_times[:click_idx + 1])
                    })

                # Track results per label
                if label_idx not in label_results:
                    label_results[label_idx] = []
                label_results[label_idx].append({
                    'image': img_path.name,
                    'dice_scores': dice_scores,
                    'final_dice': dice_scores[-1]
                })

                all_click_times.extend(click_times)

                print(f"\nSummary for {img_path.name} - Label {label_idx}:")
                print(f"  Initial Dice: {dice_scores[0]:.4f}")
                print(f"  Final Dice: {dice_scores[-1]:.4f}")
                print(f"  Improvement: {dice_scores[-1] - dice_scores[0]:.4f}")
                print(f"  Mean click time: {np.mean(click_times):.3f}s")
                print(f"  Total time: {total_time:.3f}s")

            except Exception as e:
                print(f"Error processing {img_path.name} label {label_idx}: {str(e)}")
                import traceback
                traceback.print_exc()
                continue
            finally:
                # ✅ Free GPU memory after each label
                if 'final_pred' in locals() and final_pred is not None:
                    del final_pred
                torch.cuda.empty_cache()
    # Calculate overall statistics
    if all_click_times:
        overall_mean_time = np.mean(all_click_times)
        overall_std_time = np.std(all_click_times)
        overall_median_time = np.median(all_click_times)

        # Calculate mean total time per label instance
        total_times = [results['total_time'] for results in all_results.values()]
        mean_total_time = np.mean(total_times)
        std_total_time = np.std(total_times)

        # NEW: Calculate mean time and Dice per click number across all labels
        click_number_times = {}
        click_number_dices = {}
        for click_data in all_clicks_data:
            click_num = click_data['click_number']
            if click_num not in click_number_times:
                click_number_times[click_num] = []
                click_number_dices[click_num] = []
            click_number_times[click_num].append(click_data['click_time'])
            click_number_dices[click_num].append(click_data['dice_score'])

        mean_times_per_click = {click_num: np.mean(times) for click_num, times in click_number_times.items()}
        std_times_per_click = {click_num: np.std(times) for click_num, times in click_number_times.items()}
        mean_dices_per_click = {click_num: np.mean(dices) for click_num, dices in click_number_dices.items()}
        std_dices_per_click = {click_num: np.std(dices) for click_num, dices in click_number_dices.items()}

        print(f"\n{'='*60}")
        print("OVERALL TIMING STATISTICS")
        print(f"{'='*60}")
        print(f"Mean click time: {overall_mean_time:.3f}s ± {overall_std_time:.3f}s")
        print(f"Median click time: {overall_median_time:.3f}s")
        print(f"Min click time: {np.min(all_click_times):.3f}s")
        print(f"Max click time: {np.max(all_click_times):.3f}s")
        print(f"Total clicks: {len(all_click_times)}")
        print(f"\nMean total time per label: {mean_total_time:.3f}s ± {std_total_time:.3f}s")
        print(f"Min total time: {np.min(total_times):.3f}s")
        print(f"Max total time: {np.max(total_times):.3f}s")

        # Print mean time per click number
        print(f"\n{'='*60}")
        print("MEAN TIME PER CLICK NUMBER (across all labels)")
        print(f"{'='*60}")
        for click_num in sorted(mean_times_per_click.keys()):
            print(f"Click {click_num}: {mean_times_per_click[click_num]:.3f}s ± {std_times_per_click[click_num]:.3f}s" +
                  f" (n={len(click_number_times[click_num])})")

        # Print mean Dice per click number
        print(f"\n{'='*60}")
        print("MEAN DICE PER CLICK NUMBER (across all labels)")
        print(f"{'='*60}")
        for click_num in sorted(mean_dices_per_click.keys()):
            print(f"Click {click_num}: {mean_dices_per_click[click_num]:.4f} ± {std_dices_per_click[click_num]:.4f}" +
                  f" (n={len(click_number_dices[click_num])})")

    # Print statistics per label
    print(f"\n{'='*60}")
    print("STATISTICS PER LABEL")
    print(f"{'='*60}")
    for label_idx in sorted(label_results.keys()):
        results = label_results[label_idx]
        final_dices = [r['final_dice'] for r in results]
        mean_dice = np.mean(final_dices)
        std_dice = np.std(final_dices)
        print(f"Label {label_idx}: Mean Final Dice = {mean_dice:.4f} ± {std_dice:.4f} (n={len(results)})")

    # NEW: Save results to pickle file for easy import
    results_data = {
        'all_results': all_results,
        'all_clicks_data': all_clicks_data,
        'label_results': label_results,
        'all_click_times': all_click_times,
        'statistics': {
            'overall_mean_time': overall_mean_time if all_click_times else None,
            'overall_std_time': overall_std_time if all_click_times else None,
            'overall_median_time': overall_median_time if all_click_times else None,
            'mean_total_time': mean_total_time if all_click_times else None,
            'std_total_time': std_total_time if all_click_times else None,
            'total_clicks': len(all_click_times),
            # NEW: Add mean times per click number
            'mean_times_per_click': mean_times_per_click if all_click_times else None,
            'std_times_per_click': std_times_per_click if all_click_times else None,
            'click_number_times': click_number_times if all_click_times else None,
            # NEW: Add mean Dice per click number
            'mean_dices_per_click': mean_dices_per_click if all_click_times else None,
            'std_dices_per_click': std_dices_per_click if all_click_times else None,
            'click_number_dices': click_number_dices if all_click_times else None,
        }
    }

    pickle_path = output_folder / "results_data.pkl"
    with open(pickle_path, 'wb') as f:
        pickle.dump(results_data, f)
    print(f"\nResults data saved to {pickle_path}")

    # Save summary text file
    summary_path = output_folder / "results_summary.txt"
    with open(summary_path, 'w', encoding='utf-8') as f:
        f.write("=" * 100 + "\n")
        f.write("INTERACTIVE SEGMENTATION RESULTS (MULTI-LABEL)\n")
        f.write("=" * 100 + "\n\n")

        f.write("Image\tLabel\tInitial_Dice\tFinal_Dice\tImprovement\tMean_Click_Time(s)\tTotal_Time(s)\n")
        f.write("-" * 100 + "\n")

        for result_key, results in all_results.items():
            initial = results['dice_scores'][0]
            final = results['final_dice']
            improvement = final - initial
            mean_time = results['mean_click_time']
            total_time = results['total_time']
            label = results['label']

            f.write(f"{result_key}\t{label}\t{initial:.4f}\t{final:.4f}\t{improvement:.4f}\t"
                    f"{mean_time:.3f}\t{total_time:.3f}\n")
            f.write(f"  Dice progression: {', '.join([f'{d:.4f}' for d in results['dice_scores']])}\n")
            f.write(f"  Click times: {', '.join([f'{t:.3f}s' for t in results['click_times']])}\n\n")

        if all_click_times:
            f.write("\n" + "=" * 100 + "\n")
            f.write("OVERALL TIMING STATISTICS\n")
            f.write("=" * 100 + "\n")
            f.write(f"Mean click time: {overall_mean_time:.3f}s ± {overall_std_time:.3f}s\n")
            f.write(f"Median click time: {overall_median_time:.3f}s\n")
            f.write(f"Min click time: {np.min(all_click_times):.3f}s\n")
            f.write(f"Max click time: {np.max(all_click_times):.3f}s\n")
            f.write(f"Total clicks: {len(all_click_times)}\n")
            f.write(f"\nMean total time per label: {mean_total_time:.3f}s ± {std_total_time:.3f}s\n")
            f.write(f"Min total time: {np.min(total_times):.3f}s\n")
            f.write(f"Max total time: {np.max(total_times):.3f}s\n")

            # NEW: Write mean time per click number
            f.write("\n" + "=" * 100 + "\n")
            f.write("MEAN TIME PER CLICK NUMBER (across all labels)\n")
            f.write("=" * 100 + "\n")
            for click_num in sorted(mean_times_per_click.keys()):
                f.write(
                    f"Click {click_num}: {mean_times_per_click[click_num]:.3f}s ± {std_times_per_click[click_num]:.3f}s"
                    + f" (n={len(click_number_times[click_num])})\n")

            # NEW: Write mean Dice per click number
            f.write("\n" + "=" * 100 + "\n")
            f.write("MEAN DICE PER CLICK NUMBER (across all labels)\n")
            f.write("=" * 100 + "\n")
            for click_num in sorted(mean_dices_per_click.keys()):
                f.write(
                    f"Click {click_num}: {mean_dices_per_click[click_num]:.4f} ± {std_dices_per_click[click_num]:.4f} "
                    + f"(n={len(click_number_dices[click_num])})\n")

        # Statistics per label
        f.write("\n" + "=" * 100 + "\n")
        f.write("STATISTICS PER LABEL\n")
        f.write("=" * 100 + "\n")
        for label_idx in sorted(label_results.keys()):
            results = label_results[label_idx]
            final_dices = [r['final_dice'] for r in results]
            mean_dice = np.mean(final_dices)
            std_dice = np.std(final_dices)
            f.write(f"Label {label_idx}: Mean Final Dice = {mean_dice:.4f} ± {std_dice:.4f} (n={len(results)})\n")

    print(f"\nProcessing complete! Results saved to {output_folder}")
    print(f"Summary saved to {summary_path}")

    # Plot mean Dice progression (for all labels combined)
    if all_results:
        plot_mean_dice_progression(all_results, output_folder)

        # Optionally plot per-label progression
        for label_idx in sorted(label_results.keys()):
            label_data = {r['image']: {'dice_scores': r['dice_scores']} for r in label_results[label_idx]}
            label_output_folder = output_folder / f"label_{label_idx}"
            label_output_folder.mkdir(exist_ok=True)
            plot_mean_dice_progression(label_data, label_output_folder)

    return results_data


# --- Main Execution ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process medical images for interactive segmentation evaluation')

    # Required arguments
    parser.add_argument('--images_dir',
                        type=str,
                        required=True,
                        help='Path to folder containing input images (NIfTI format: .nii or .nii.gz)')

    parser.add_argument('--labels_dir',
                        type=str,
                        required=True,
                        help='Path to folder containing ground truth label masks')

    parser.add_argument('--work_dir', type=str, required=True, help='Path to folder where results will be saved')

    # Optional arguments
    parser.add_argument('--click_number',
                        type=int,
                        default=50,
                        help='Number of slices to sample clicks from (default: 4)')

    parser.add_argument(
        '--labels',
        type=int,
        nargs='+',
        default=None,
        help='Specific label indices to process (default: all non-zero labels). Example: --labels 1 2 3')

    parser.add_argument("--checkpoint", type=str, help="Path to the model checkpoint file")

    args = parser.parse_args()
    # Create the trainer/inference module
    predict = Inference_module(
        device="cuda",
        checkpoint=args.checkpoint,
        activate_propagation=True,  # Enable automatic propagation
        torch_compile=False,
        vis=False,
    )

    print("Configuration:")
    print(f"  Images folder: {args.images_dir}")
    print(f"  Ground truth folder: {args.labels_dir}")
    print(f"  Output folder: {args.work_dir}")
    if args.labels:
        print(f"  Processing labels: {args.labels}")
    else:
        print("  Processing labels: all non-zero labels")
    print()

    process_folder(image_folder=args.images_dir,
                   gt_folder=args.labels_dir,
                   output_folder=args.work_dir,
                   click_number=args.click_number,
                   labels_to_process=args.labels)