example_usage.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)
"""
Example usage of the inference module.
It is recomended to have at least 10Go of Vram for inference with propagation

This script demonstrates how to:
1. Initialize the model
2. Process an image and ground truth label
3. Perform inference with user clicks
"""

from pathlib import Path

import torch

from SLIP_model.inference_module import Inference_module


def main():
    # ========================================
    # CONFIGURATION
    # ========================================

    # Paths to your data in .nii.gz
    image_path = Path("path_to_image.nii.gz")  # Change this to your image path
    label_path = Path("path_to_label.nii.gz")  # Change this to your label path

    # Model checkpoint
    checkpoint_path = "path_to_ckpt.pth"

    # Device configuration
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Target label to segment (set to None to auto-select first non-zero label)
    target_label = 1

    # ========================================
    # STEP 1: Initialize the model
    # ========================================
    print("=" * 60)
    print("Initializing model...")
    print("=" * 60)

    # Create the trainer/inference module
    predict = Inference_module(
        device=device,
        checkpoint=checkpoint_path,
        activate_propagation=True,  # Enable automatic propagation
        torch_compile=True,
        vis=True,
    )

    print(f"Model loaded on device: {device}")

    # ========================================
    # STEP 2: Process the image
    # ========================================
    print("\n" + "=" * 60)
    print("Processing image and computing embeddings...")
    print("=" * 60)

    predict.process_image(
        image_path=image_path,
        gt_path=label_path,  # Could be set to None for inference
        target_label=target_label,
    )

    print("Image processing complete!")
    print(f"Volume shape: {predict.vol.shape}")
    print(f"Ground truth shape: {predict.gt3D.shape}")
    print(f"Number of patches: {len(predict.bbox_data)}")

    # ========================================
    # STEP 3: Perform inference with clicks
    # ========================================
    print("\n" + "=" * 60)
    print("Performing inference with user clicks...")
    print("=" * 60)

    # ── Click 1 ──────────────────────────────────────────────────────────────
    print("\n--- Click 1 ---")
    clicks_1 = [[105, 119, 79]]
    labels_1 = [1]  # 1 for positive click
    pred_masks_1 = predict.click_inference(clicks_1, labels_1)

    # ── Click 2 ──────────────────────────────────────────────────────────────
    print("\n--- Click 2 ---")
    clicks_2 = [[106, 119, 75]]
    labels_2 = [0]  # 0 for negative click
    pred_masks_2 = predict.click_inference(clicks_2, labels_2)

    # ── Undo (revert to state after click 1) ─────────────────────────────────
    print("\n--- Undo last click ---")
    predict.undo(predict.action_history)
    print("Undo complete.")

    # ========================================
    # RESULTS
    # ========================================
    print("\n" + "=" * 60)
    print("RESULTS")
    print("=" * 60)
    print(f"Prediction shape   : {pred_masks_2.shape}")
    pred_masks = pred_masks_2

    # You can save the prediction if needed
    # import SimpleITK as sitk
    # pred_binary = (pred_masks.squeeze().cpu().numpy() > 0).astype(np.uint8)
    # pred_image = sitk.GetImageFromArray(pred_binary)
    # sitk.WriteImage(pred_image, "prediction.nii.gz")
    # print("Prediction saved to prediction.nii.gz")

    return pred_masks


if __name__ == "__main__":
    # Check if CUDA is available
    if not torch.cuda.is_available():
        print("WARNING: CUDA is not available. Running on CPU")
        print("Consider using a GPU for inference.")

    # Run the main example
    try:
        main()
    except Exception as e:
        print(f"\nError occurred: {e}")
        print("\nPlease make sure to:")
        print("1. Update the image_path and label_path variables")
        print("2. Ensure the checkpoint file exists")
        print("3. Have sufficient GPU memory")
        import traceback
        traceback.print_exc()
        traceback.print_exc()
        traceback.print_exc()