Upload /Interactive_MEN_RT_predictor.py with huggingface_hub

Interactive_MEN_RT_predictor.py

@@ -0,0 +1,1012 @@
+import os
+import numpy as np
+import torch
+from typing import Union, List, Tuple, Optional, Dict
+from concurrent.futures import ThreadPoolExecutor
+from time import time
+import sys
+import importlib
+import math
+
+from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
+from nnunetv2.utilities.helpers import empty_cache, dummy_context
+from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels
+from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager
+from batchgenerators.utilities.file_and_folder_operations import load_json, join, subdirs
+from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice, crop_and_pad_nd
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.functional import interpolate
+import nnunetv2
+
+import nnInteractive
+from nnInteractive.interaction.point import PointInteraction_stub
+from nnInteractive.utils.bboxes import generate_bounding_boxes
+from nnInteractive.utils.crop import crop_and_pad_into_buffer, paste_tensor, pad_cropped, crop_to_valid
+from nnInteractive.utils.erosion_dilation import iterative_3x3_same_padding_pool3d
+from nnInteractive.utils.rounding import round_to_nearest_odd
+
+
+class InteractiveMENRTPredictor:
+    """
+    Interactive MEN RT Predictor for interactive segmentation with point, bbox, scribble, and lasso interactions.
+    """
+    
+    def __init__(self,
+                 device: torch.device = torch.device('cuda'),
+                 use_torch_compile: bool = False,
+                 verbose: bool = False,
+                 torch_n_threads: int = 8,
+                 do_autozoom: bool = True,
+                 use_pinned_memory: bool = True
+    ):
+        """
+        Only intended to work with nnInteractiveTrainerV2 and its derivatives
+        """
+        # set as part of initialization
+        assert use_torch_compile is False, ('This implementation places the preprocessed image and the interactions '
+                                            'into pinned memory for speed reasons. This is incompatible with '
+                                            'torch.compile because of inconsistent strides in the memory layout. '
+                                            'Note to self: .contiguous() on GPU could be a solution. Unclear whether '
+                                            'that will yield a benefit though.')
+        self.network = None
+        self.label_manager = None
+        self.dataset_json = None
+        self.trainer_name = None
+        self.configuration_manager = None
+        self.plans_manager = None
+        self.use_pinned_memory = use_pinned_memory
+        self.device = device
+        self.use_torch_compile = use_torch_compile
+        
+        # Interactive session state
+        self.interactions: torch.Tensor = None
+        self.preprocessed_image: torch.Tensor = None
+        self.preprocessed_props = None
+        self.target_buffer: Union[np.ndarray, torch.Tensor] = None
+        
+        self.pad_mode_data = self.preferred_scribble_thickness = self.point_interaction = None
+        self.verbose = verbose
+        
+        self.do_autozoom: bool = do_autozoom
+        torch.set_num_threads(min(torch_n_threads, os.cpu_count()))
+
+        self.original_image_shape = None
+
+        self.new_interaction_zoom_out_factors: List[float] = []
+        self.new_interaction_centers = []
+        self.has_positive_bbox = False
+
+        # Create a thread pool executor for background tasks.
+        # this only takes care of preprocessing and interaction memory initialization so there is no need to give it
+        # more than 2 workers
+        self.executor = ThreadPoolExecutor(max_workers=2)
+        self.preprocess_future = None
+        self.interactions_future = None
+
+    def set_image(self, image: np.ndarray, image_properties: dict = None):
+        """
+        Image must be 4D to satisfy nnU-Net needs: [c, x, y, z]
+        Offload the processing to a background thread.
+        """
+        if image_properties is None:
+            image_properties = {}
+        self._reset_session()
+        assert image.ndim == 4, f'expected a 4d image as input, got {image.ndim}d. Shape {image.shape}'
+        if self.verbose:
+            print(f'Initialize with raw image shape {image.shape}')
+
+        # Offload all image preprocessing to a background thread.
+        self.preprocess_future = self.executor.submit(self._background_set_image, image, image_properties)
+        self.original_image_shape = image.shape
+
+    def _finish_preprocessing_and_initialize_interactions(self):
+        """
+        Block until both the image preprocessing and the interactions tensor initialization
+        are finished.
+        """
+        if self.preprocess_future is not None:
+            # Wait for image preprocessing to complete.
+            self.preprocess_future.result()
+            del self.preprocess_future
+            self.preprocess_future = None
+
+    def set_target_buffer(self, target_buffer: Union[np.ndarray, torch.Tensor]):
+        """
+        Must be 3d numpy array or torch.Tensor
+        """
+        self.target_buffer = target_buffer
+
+    def set_do_autozoom(self, do_propagation: bool, max_num_patches: Optional[int] = None):
+        self.do_autozoom = do_propagation
+
+    def _reset_session(self):
+        self.interactions_future = None
+        self.preprocess_future = None
+
+        del self.preprocessed_image
+        del self.target_buffer
+        del self.interactions
+        del self.preprocessed_props
+        self.preprocessed_image = None
+        self.target_buffer = None
+        self.interactions = None
+        self.preprocessed_props = None
+        empty_cache(self.device)
+        self.original_image_shape = None
+        self.has_positive_bbox = False
+
+    def _initialize_interactions(self, image_torch: torch.Tensor):
+        if self.verbose:
+            print(f'Initialize interactions. Pinned: {self.use_pinned_memory}')
+        # Create the interaction tensor based on the target shape.
+        self.interactions = torch.zeros(
+            (7, *image_torch.shape[1:]),
+            device='cpu',
+            dtype=torch.float16,
+            pin_memory=(self.device.type == 'cuda' and self.use_pinned_memory)
+        )
+    
+    def _background_set_image(self, image: np.ndarray, image_properties: dict):
+        """Background preprocessing of the image"""
+        # Convert to torch tensor
+        image_torch = torch.clone(torch.from_numpy(image))
+        
+        # Crop to nonzero region
+        if self.verbose:
+            print('Cropping input image to nonzero region')
+        nonzero_idx = torch.where(image_torch != 0)
+        bbox = [[i.min().item(), i.max().item() + 1] for i in nonzero_idx]
+        
+        # Ensure bbox is larger than patch_size
+        if hasattr(self, 'configuration_manager') and self.configuration_manager is not None:
+            patch_size = self.configuration_manager.patch_size
+            for dim in range(1, len(bbox)):
+                bbox_size = bbox[dim][1] - bbox[dim][0]
+                if bbox_size < patch_size[dim - 1]:
+                    # Center the bbox and extend it to patch_size
+                    center = (bbox[dim][0] + bbox[dim][1]) // 2
+                    bbox[dim][0] = max(0, center - patch_size[dim - 1] // 2)
+                    bbox[dim][1] = min(image_torch.shape[dim], center + patch_size[dim - 1] // 2 + patch_size[dim - 1] % 2)
+        
+        del nonzero_idx
+        slicer = bounding_box_to_slice(bbox)
+        image_torch = image_torch[slicer].float()
+        
+        if self.verbose:
+            print(f'Cropped image shape: {image_torch.shape}')
+            
+        # Initialize interactions tensor
+        self._initialize_interactions(image_torch)
+        
+        # Normalize the image
+        if self.verbose:
+            print('Normalizing cropped image')
+        image_torch -= image_torch.mean()
+        image_torch /= image_torch.std()
+        
+        self.preprocessed_image = image_torch
+        if self.use_pinned_memory and self.device.type == 'cuda':
+            if self.verbose:
+                print('Pin memory: image')
+            self.preprocessed_image = self.preprocessed_image.pin_memory()
+            
+        self.preprocessed_props = {'bbox_used_for_cropping': bbox[1:]}
+
+    def reset_interactions(self):
+        """
+        Use this to reset all interactions and start from scratch for the current image. This includes the initial
+        segmentation!
+        """
+        if self.interactions is not None:
+            self.interactions.fill_(0)
+
+        if self.target_buffer is not None:
+            if isinstance(self.target_buffer, np.ndarray):
+                self.target_buffer.fill(0)
+            elif isinstance(self.target_buffer, torch.Tensor):
+                self.target_buffer.zero_()
+        empty_cache(self.device)
+        self.has_positive_bbox = False
+    
+    def add_bbox_interaction(self, bbox_coords, include_interaction: bool, run_prediction: bool = True) -> np.ndarray:
+        if include_interaction:
+            self.has_positive_bbox = True
+
+        self._finish_preprocessing_and_initialize_interactions()
+
+        lbs_transformed = [round(i) for i in transform_coordinates_noresampling([i[0] for i in bbox_coords],
+                                                             self.preprocessed_props['bbox_used_for_cropping'])]
+        ubs_transformed = [round(i) for i in transform_coordinates_noresampling([i[1] for i in bbox_coords],
+                                                             self.preprocessed_props['bbox_used_for_cropping'])]
+        transformed_bbox_coordinates = [[i, j] for i, j in zip(lbs_transformed, ubs_transformed)]
+
+        if self.verbose:
+            print(f'Added bounding box coordinates.\n'
+                  f'Raw: {bbox_coords}\n'
+                  f'Transformed: {transformed_bbox_coordinates}\n'
+                  f"Crop Bbox: {self.preprocessed_props['bbox_used_for_cropping']}")
+
+        # Prevent collapsed bounding boxes and clip to image shape
+        image_shape = self.preprocessed_image.shape  # Assuming shape is (C, H, W, D) or similar
+
+        for dim in range(len(transformed_bbox_coordinates)):
+            transformed_start, transformed_end = transformed_bbox_coordinates[dim]
+
+            # Clip to image boundaries
+            transformed_start = max(0, transformed_start)
+            transformed_end = min(image_shape[dim + 1], transformed_end)  # +1 to skip channel dim
+
+            # Ensure the bounding box does not collapse to a single point
+            if transformed_end <= transformed_start:
+                if transformed_start == 0:
+                    transformed_end = min(1, image_shape[dim + 1])
+                else:
+                    transformed_start = max(transformed_start - 1, 0)
+
+            transformed_bbox_coordinates[dim] = [transformed_start, transformed_end]
+
+        if self.verbose:
+            print(f'Bbox coordinates after clip to image boundaries and preventing dim collapse:\n'
+                  f'Bbox: {transformed_bbox_coordinates}\n'
+                  f'Internal image shape: {self.preprocessed_image.shape}')
+
+        self._add_patch_for_bbox_interaction(transformed_bbox_coordinates)
+
+        # decay old interactions
+        self.interactions[-6:-4] *= self.interaction_decay
+
+        # place bbox
+        slicer = tuple([slice(*i) for i in transformed_bbox_coordinates])
+        channel = -6 if include_interaction else -5
+        self.interactions[(channel, *slicer)] = 1
+
+        # forward pass
+        if run_prediction:
+            self._predict()
+
+    def add_point_interaction(self, coordinates: Tuple[int, ...], include_interaction: bool, run_prediction: bool = True):
+        self._finish_preprocessing_and_initialize_interactions()
+
+        transformed_coordinates = [round(i) for i in transform_coordinates_noresampling(coordinates,
+                                                             self.preprocessed_props['bbox_used_for_cropping'])]
+
+        self._add_patch_for_point_interaction(transformed_coordinates)
+
+        # decay old interactions
+        self.interactions[-4:-2] *= self.interaction_decay
+
+        interaction_channel = -4 if include_interaction else -3
+        self.interactions[interaction_channel] = self.point_interaction.place_point(
+            transformed_coordinates, self.interactions[interaction_channel])
+        if run_prediction:
+            self._predict()
+
+    def add_scribble_interaction(self, scribble_image: np.ndarray,  include_interaction: bool, run_prediction: bool = True):
+        assert all([i == j for i, j in zip(self.original_image_shape[1:], scribble_image.shape)]), f'Given scribble image must match input image shape. Input image was: {self.original_image_shape[1:]}, given: {scribble_image.shape}'
+        self._finish_preprocessing_and_initialize_interactions()
+
+        scribble_image = torch.from_numpy(scribble_image)
+
+        # crop (as in preprocessing)
+        scribble_image = crop_and_pad_nd(scribble_image, self.preprocessed_props['bbox_used_for_cropping'])
+
+        self._add_patch_for_scribble_interaction(scribble_image)
+
+        # decay old interactions
+        self.interactions[-2:] *= self.interaction_decay
+
+        interaction_channel = -2 if include_interaction else -1
+        torch.maximum(self.interactions[interaction_channel], scribble_image.to(self.interactions.device),
+                      out=self.interactions[interaction_channel])
+        del scribble_image
+        empty_cache(self.device)
+        if run_prediction:
+            self._predict()
+
+    def add_lasso_interaction(self, lasso_image: np.ndarray,  include_interaction: bool, run_prediction: bool = True):
+        assert all([i == j for i, j in zip(self.original_image_shape[1:], lasso_image.shape)]), f'Given lasso image must match input image shape. Input image was: {self.original_image_shape[1:]}, given: {lasso_image.shape}'
+        self._finish_preprocessing_and_initialize_interactions()
+
+        lasso_image = torch.from_numpy(lasso_image)
+
+        # crop (as in preprocessing)
+        lasso_image = crop_and_pad_nd(lasso_image, self.preprocessed_props['bbox_used_for_cropping'])
+
+        self._add_patch_for_lasso_interaction(lasso_image)
+
+        # decay old interactions
+        self.interactions[-6:-4] *= self.interaction_decay
+
+        # lasso is written into bbox channel
+        interaction_channel = -6 if include_interaction else -5
+        torch.maximum(self.interactions[interaction_channel], lasso_image.to(self.interactions.device),
+                      out=self.interactions[interaction_channel])
+        del lasso_image
+        empty_cache(self.device)
+        if run_prediction:
+            self._predict()
+
+    def add_initial_seg_interaction(self, initial_seg: np.ndarray, run_prediction: bool = False):
+        """
+        WARNING THIS WILL RESET INTERACTIONS!
+        """
+        assert all([i == j for i, j in zip(self.original_image_shape[1:], initial_seg.shape)]), f'Given initial seg must match input image shape. Input image was: {self.original_image_shape[1:]}, given: {initial_seg.shape}'
+
+        self._finish_preprocessing_and_initialize_interactions()
+
+        self.reset_interactions()
+
+        if isinstance(self.target_buffer, np.ndarray):
+            self.target_buffer[:] = initial_seg
+
+        initial_seg = torch.from_numpy(initial_seg)
+
+        if isinstance(self.target_buffer, torch.Tensor):
+            self.target_buffer[:] = initial_seg
+
+        # crop (as in preprocessing)
+        initial_seg = crop_and_pad_nd(initial_seg, self.preprocessed_props['bbox_used_for_cropping'])
+
+        # initial seg is written into initial seg buffer
+        interaction_channel = -7
+        self.interactions[interaction_channel] = initial_seg
+        empty_cache(self.device)
+        if run_prediction:
+            self._add_patch_for_initial_seg_interaction(initial_seg)
+            del initial_seg
+            self._predict()
+        else:
+            del initial_seg
+
+    @torch.inference_mode()
+    def _predict(self):
+        """
+        Perform prediction with interactions. The process follows the training procedure:
+        1. Make initial prediction with current interactions
+        2. Generate new interactions based on prediction errors
+        3. Make final prediction with updated interactions
+        """
+        assert self.pad_mode_data == 'constant', 'pad modes other than constant are not implemented here'
+
+        start_predict = time()
+        with torch.autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
+            # Find the region containing all interactions
+            interaction_mask = torch.any(self.interactions[1:] > 0, dim=0)  # Combine all interaction channels
+            if not torch.any(interaction_mask):
+                print('No interactions found, skipping prediction')
+                return
+
+            # Get bounding box of interaction region
+            nonzero_indices = torch.nonzero(interaction_mask)
+            min_coords = torch.min(nonzero_indices, dim=0)[0]
+            max_coords = torch.max(nonzero_indices, dim=0)[0]
+            
+            # Initialize bbox with interaction region
+            patch_size = self.configuration_manager.patch_size
+            half_patch_size = [p // 2 for p in patch_size]
+            image_shape = self.preprocessed_image.shape[1:]
+            
+            # For each dimension, calculate bbox ensuring:
+            # 1. bbox start >= 0
+            # 2. bbox end <= image_shape
+            # 3. bbox size >= patch_size
+            bbox = []
+            for i, (min_c, max_c, h, p) in enumerate(zip(min_coords, max_coords, half_patch_size, patch_size)):
+                start = max(0, min(image_shape[i] - p, (min_c + max_c) // 2 - p // 2))
+                end = min(image_shape[i], start + p)
+                bbox.append([start, end])
+
+            # Calculate number of patches needed
+            overlap = [64, 64, 64]                             # [O_z, O_y, O_x]
+            num_patches = [
+                1 if (b1 - b0) <= P 
+                else math.ceil(((b1 - b0) - P) / (P - O)) + 1
+                for (b0, b1), P, O in zip(bbox, patch_size, overlap)
+            ]
+            
+            # Initialize prediction tensors for soft merging
+            final_pred_soft = torch.zeros((2, *self.preprocessed_image.shape[1:]), dtype=torch.float32, device='cpu')
+            prediction_count = torch.zeros(self.preprocessed_image.shape[1:], dtype=torch.float32, device='cpu')
+
+            # Process each patch
+            for x in range(num_patches[0]):
+                for y in range(num_patches[1]):
+                    for z in range(num_patches[2]):
+                        # Calculate patch boundaries
+                        step_index = [x, y, z]
+                        start_coords = [bbox[i][0] + step_index[i] * p for i, p in zip([0, 1, 2], patch_size)]
+                        end_coords = [min(bbox[i][1], start_coords[i] + p) for i, p in zip([0, 1, 2], patch_size)]
+                        
+                        for i in range(len(patch_size)):
+                            if end_coords[i] - start_coords[i] < patch_size[i]:
+                                if end_coords[i] >= bbox[i][1]:
+                                    start_coords[i] = bbox[i][1] - patch_size[i]
+                        
+                        # Extract image patch
+                        image_patch = self.preprocessed_image[:, start_coords[0]:end_coords[0], 
+                                                           start_coords[1]:end_coords[1], 
+                                                           start_coords[2]:end_coords[2]]
+                        
+                        # Extract interaction patches
+                        interaction_patch = self.interactions[:, start_coords[0]:end_coords[0],
+                                                           start_coords[1]:end_coords[1],
+                                                           start_coords[2]:end_coords[2]]
+                        
+                        # Pad to patch_size if necessary
+                        if not all([e - s == p for s, e, p in zip(start_coords, end_coords, patch_size)]):
+                            pad_size = [(0, p - (e - s)) for s, e, p in zip(start_coords, end_coords, patch_size)]
+                            image_patch = F.pad(image_patch, [item for sublist in reversed(pad_size) for item in sublist])
+                            interaction_patch = F.pad(interaction_patch, [item for sublist in reversed(pad_size) for item in sublist])
+                        
+                        # Move to device
+                        image_patch = image_patch.to(self.device, non_blocking=self.device.type == 'cuda')
+                        interaction_patch = interaction_patch.to(self.device, non_blocking=self.device.type == 'cuda')
+                        
+                        # Concatenate image and interaction channels
+                        input_for_predict = torch.cat((image_patch, interaction_patch))
+                        
+                        # Make prediction
+                        pred_raw = self.network(input_for_predict[None])[0]
+                        pred_prob = F.softmax(pred_raw, dim=0)
+                        
+                        del input_for_predict, pred_raw, image_patch, interaction_patch
+                        
+                        # Resize prediction if needed
+                        if not all([e - s == p for s, e, p in zip(start_coords, end_coords, patch_size)]):
+                            pred_prob = interpolate(pred_prob[None], 
+                                                  [e - s for s, e in zip(start_coords, end_coords)], 
+                                                  mode='trilinear')[0]
+
+                        # Add to accumulated predictions
+                        pred_prob = pred_prob.cpu()
+                        final_pred_soft[:, start_coords[0]:end_coords[0],
+                                      start_coords[1]:end_coords[1],
+                                      start_coords[2]:end_coords[2]] += pred_prob
+                        prediction_count[start_coords[0]:end_coords[0],
+                                       start_coords[1]:end_coords[1],
+                                       start_coords[2]:end_coords[2]] += 1
+
+                        del pred_prob
+                        empty_cache(self.device)
+
+            # Average predictions and convert to binary
+            final_pred_soft = final_pred_soft / prediction_count.clamp(min=1)
+            # final_pred_soft = self._iterative_adjust_prediction(final_pred_soft, self.interactions)
+            final_pred = (final_pred_soft[1] >= 0.5).to(torch.uint8)
+
+            # Update interactions and target buffer
+            self.interactions[0][:] = final_pred
+            paste_tensor(self.target_buffer, final_pred, self.preprocessed_props['bbox_used_for_cropping'])
+
+        print(f'Done. Total time {round(time() - start_predict, 3)}s')
+
+        self.new_interaction_centers = []
+        empty_cache(self.device)
+
+    @torch.inference_mode()
+    def _predict_without_interaction(self):
+        """
+        Perform prediction with interaction channels but without zooming. This is a simplified version of _predict that:
+        1. Makes prediction on the entire image at once using interaction channels
+        2. No zooming or refinement is performed
+        3. Uses all interaction channels (previous segmentation, bbox, point, scribble)
+        """
+        assert self.pad_mode_data == 'constant', 'pad modes other than constant are not implemented here'
+        
+        start_predict = time()
+        with torch.autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
+            # Get image dimensions
+            image_shape = self.preprocessed_image.shape[1:]  # Remove channel dimension
+            
+            # Calculate number of patches needed
+            patch_size = self.configuration_manager.patch_size
+            bbox = [[0, i] for i in image_shape]
+
+            # Calculate number of patches needed
+            overlap = [64, 64, 64]                             # [O_z, O_y, O_x]
+            num_patches = [
+                1 if (b1 - b0) <= P 
+                else math.ceil(((b1 - b0) - P) / (P - O)) + 1
+                for (b0, b1), P, O in zip(bbox, patch_size, overlap)
+            ]
+            
+            # Initialize prediction tensors for soft merging
+            pred_soft = torch.zeros((2, *image_shape), dtype=torch.float32, device='cpu')  # 2 channels for binary segmentation
+            pred_count = torch.zeros(image_shape, dtype=torch.float32, device='cpu')
+            
+            # Process each patch
+            for x in range(num_patches[0]):
+                for y in range(num_patches[1]):
+                    for z in range(num_patches[2]):
+                        # Calculate patch boundaries
+                        step_index = [x, y, z]
+                        start_coords = [bbox[i][0] + step_index[i] * p for i, p in zip([0, 1, 2], patch_size)]
+                        end_coords = [min(bbox[i][1], start_coords[i] + p) for i, p in zip([0, 1, 2], patch_size)]
+                        
+                        for i in range(len(patch_size)):
+                            if end_coords[i] - start_coords[i] < patch_size[i]:
+                                if end_coords[i] >= bbox[i][1]:
+                                    start_coords[i] = bbox[i][1] - patch_size[i]
+                        
+                        # Extract image patch
+                        image_patch = self.preprocessed_image[:, start_coords[0]:end_coords[0], 
+                                                           start_coords[1]:end_coords[1], 
+                                                           start_coords[2]:end_coords[2]]
+                        
+                        # Extract interaction patches
+                        interaction_patch = self.interactions[:, start_coords[0]:end_coords[0],
+                                                           start_coords[1]:end_coords[1],
+                                                           start_coords[2]:end_coords[2]]
+                        
+                        # Pad if necessary
+                        if not all([e - s == p for s, e, p in zip(start_coords, end_coords, patch_size)]):
+                            pad_size = [(0, p - (e - s)) for s, e, p in zip(start_coords, end_coords, patch_size)]
+                            image_patch = F.pad(image_patch, [item for sublist in reversed(pad_size) for item in sublist])
+                            interaction_patch = F.pad(interaction_patch, [item for sublist in reversed(pad_size) for item in sublist])
+                        
+                        # Move to device
+                        image_patch = image_patch.to(self.device, non_blocking=self.device.type == 'cuda')
+                        interaction_patch = interaction_patch.to(self.device, non_blocking=self.device.type == 'cuda')
+                        
+                        # Concatenate image and interaction channels
+                        input_for_predict = torch.cat((image_patch, interaction_patch))
+                        
+                        # Make prediction and get soft probabilities
+                        patch_pred = self.network(input_for_predict[None])[0]
+                        patch_prob = F.softmax(patch_pred, dim=0)
+                        
+                        # Resize prediction to original patch size if necessary
+                        if not all([e - s == p for s, e, p in zip(start_coords, end_coords, patch_size)]):
+                            patch_prob = interpolate(patch_prob[None], 
+                                                   [e - s for s, e in zip(start_coords, end_coords)], 
+                                                   mode='trilinear')[0]
+                        
+                        # Add to accumulated predictions
+                        pred_soft[:, start_coords[0]:end_coords[0], 
+                                start_coords[1]:end_coords[1], 
+                                start_coords[2]:end_coords[2]] += patch_prob.cpu()
+                        pred_count[start_coords[0]:end_coords[0], 
+                                 start_coords[1]:end_coords[1], 
+                                 start_coords[2]:end_coords[2]] += 1
+                        
+                        del image_patch, interaction_patch, input_for_predict, patch_pred, patch_prob
+                        empty_cache(self.device)
+            
+            # Average predictions and convert to binary
+            pred_soft = pred_soft / pred_count.clamp(min=1)
+            pred = (pred_soft[1] >= 0.5).to(torch.uint8)
+
+            # Update interactions and target buffer
+            self.interactions[0][:] = pred
+            paste_tensor(self.target_buffer, pred, self.preprocessed_props['bbox_used_for_cropping'])
+        
+        print(f'Done. Total time {round(time() - start_predict, 3)}s')
+        empty_cache(self.device)
+
+    def _add_patch_for_point_interaction(self, coordinates):
+        self.new_interaction_centers.append(coordinates)
+        print(f'Added new point interaction: center {coordinates}')
+
+    def _add_patch_for_bbox_interaction(self, bbox):
+        bbox_center = [round((i[0] + i[1]) / 2) for i in bbox]
+        bbox_size = [i[1]-i[0] for i in bbox]
+        # we want to see some context, so the crop we see for the initial prediction should be patch_size / 3 larger
+        requested_size = [i + j // 3 for i, j in zip(bbox_size, self.configuration_manager.patch_size)]
+        self.new_interaction_zoom_out_factors.append(max(1, max([i / j for i, j in zip(requested_size, self.configuration_manager.patch_size)])))
+        self.new_interaction_centers.append(bbox_center)
+        print(f'Added new bbox interaction: center {bbox_center}')
+
+    def _add_patch_for_scribble_interaction(self, scribble_image):
+        return self._generic_add_patch_from_image(scribble_image)
+
+    def _add_patch_for_lasso_interaction(self, lasso_image):
+        return self._generic_add_patch_from_image(lasso_image)
+
+    def _add_patch_for_initial_seg_interaction(self, initial_seg):
+        return self._generic_add_patch_from_image(initial_seg)
+
+    def _generic_add_patch_from_image(self, image: torch.Tensor):
+        if not torch.any(image):
+            print('Received empty image prompt. Cannot add patches for prediction')
+            return
+        nonzero_indices = torch.nonzero(image, as_tuple=False)
+        mn = torch.min(nonzero_indices, dim=0)[0]
+        mx = torch.max(nonzero_indices, dim=0)[0]
+        roi = [[i.item(), x.item() + 1] for i, x in zip(mn, mx)]
+        roi_center = [round((i[0] + i[1]) / 2) for i in roi]
+        roi_size = [i[1]- i[0] for i in roi]
+        requested_size = [i + j // 3 for i, j in zip(roi_size, self.configuration_manager.patch_size)]
+        self.new_interaction_zoom_out_factors.append(max(1, max([i / j for i, j in zip(requested_size, self.configuration_manager.patch_size)])))
+        self.new_interaction_centers.append(roi_center)
+        print(f'Added new image interaction: scale {self.new_interaction_zoom_out_factors[-1]}, center {roi_center}')
+        
+    def initialize_from_trained_model_folder(self, 
+                                           model_training_output_dir: str,
+                                           use_fold: Union[int, str] = None,
+                                           checkpoint_name: str = 'checkpoint_final.pth'):
+        """
+        Initialize the predictor from a trained model folder.
+        """
+        # Determine fold folder
+        if use_fold is not None:
+            use_fold = int(use_fold) if use_fold != 'all' else use_fold
+            fold_folder = f'fold_{use_fold}'
+        else:
+            fldrs = subdirs(model_training_output_dir, prefix='fold_', join=False)
+            assert len(fldrs) == 1, f'Attempted to infer fold but there is != 1 fold_ folders: {fldrs}'
+            fold_folder = fldrs[0]
+            
+        # load trainer specific settings
+        expected_json_file = join(model_training_output_dir, fold_folder, 'inference_session_class.json')
+        json_content = load_json(expected_json_file)
+        if isinstance(json_content, str):
+            # Old convention where we only specified the inference class in this file
+            point_interaction_radius = 4
+            point_interaction_use_etd = True
+            self.preferred_scribble_thickness = [2, 2, 2]
+            self.point_interaction = PointInteraction_stub(
+                point_interaction_radius,
+                point_interaction_use_etd)
+            self.pad_mode_data = "constant"
+            self.interaction_decay = 0.9
+        else:
+            point_interaction_radius = json_content['point_radius']
+            self.preferred_scribble_thickness = json_content['preferred_scribble_thickness']
+            if not isinstance(self.preferred_scribble_thickness, (tuple, list)):
+                self.preferred_scribble_thickness = [self.preferred_scribble_thickness] * 3
+            self.interaction_decay = json_content['interaction_decay'] if 'interaction_decay' in json_content.keys() else 0.9
+            point_interaction_use_etd = json_content['use_distance_transform'] if 'use_distance_transform' in json_content.keys() else True
+            self.point_interaction = PointInteraction_stub(point_interaction_radius, point_interaction_use_etd)
+            # padding mode for data. See nnInteractiveTrainerV2_nodelete_reflectpad
+            self.pad_mode_data = json_content['pad_mode_image'] if 'pad_mode_image' in json_content.keys() else "constant"
+
+        # Load dataset and plans
+        dataset_json = load_json(join(model_training_output_dir, 'dataset.json'))
+        plans = load_json(join(model_training_output_dir, 'plans.json'))
+        plans_manager = PlansManager(plans)
+
+        # Load checkpoint
+        checkpoint = torch.load(join(model_training_output_dir, fold_folder, checkpoint_name),
+                              map_location=self.device, weights_only=False)
+        trainer_name = checkpoint['trainer_name']
+        configuration_name = checkpoint['init_args']['configuration']
+        parameters = checkpoint['network_weights']
+
+        # Get configuration
+        configuration_manager = plans_manager.get_configuration(configuration_name)
+        
+        # Restore network
+        num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json)
+        network = nnUNetTrainer.build_network_architecture(
+            configuration_manager.network_arch_class_name,
+            configuration_manager.network_arch_init_kwargs,
+            configuration_manager.network_arch_init_kwargs_req_import,
+            num_input_channels,
+            plans_manager.get_label_manager(dataset_json).num_segmentation_heads,
+            enable_deep_supervision=False
+        ).to(self.device)
+        network.load_state_dict(parameters)
+
+        # Store necessary information
+        self.plans_manager = plans_manager
+        self.configuration_manager = configuration_manager
+        self.network = network
+        self.dataset_json = dataset_json
+        self.trainer_name = trainer_name
+        self.label_manager = plans_manager.get_label_manager(dataset_json)
+
+        if self.use_torch_compile:
+            print('Using torch.compile')
+            self.network = torch.compile(self.network)
+
+        if self.verbose:
+            print(f"Loaded interactive config: point_radius={self.point_interaction.point_radius}, "
+                  f"scribble_thickness={self.preferred_scribble_thickness}, "
+                  f"interaction_decay={self.interaction_decay}")
+            
+    def manual_initialization(self, network: nn.Module, plans_manager: PlansManager,
+                              configuration_manager: ConfigurationManager,
+                              dataset_json: dict, trainer_name: str):
+        """
+        This is used by the nnUNetTrainer to initialize nnUNetPredictor for the final validation
+        """
+        self.plans_manager = plans_manager
+        self.configuration_manager = configuration_manager
+        self.network = network
+        self.dataset_json = dataset_json
+        self.trainer_name = trainer_name
+        self.label_manager = plans_manager.get_label_manager(dataset_json)
+
+        if self.use_torch_compile and not isinstance(self.network, OptimizedModule):
+            print('Using torch.compile')
+            self.network = torch.compile(self.network)
+
+        if not self.use_torch_compile and isinstance(self.network, OptimizedModule):
+            self.network = self.network._orig_mod
+
+        self.network = self.network.to(self.device)
+
+    @torch.inference_mode()
+    def _predict_autozoom(self):
+        """
+        Perform prediction with interactions. The process follows the training procedure:
+        1. Make initial prediction with current interactions
+        2. Generate new interactions based on prediction errors
+        3. Make final prediction with updated interactions
+        """
+        assert self.pad_mode_data == 'constant', 'pad modes other than constant are not implemented here'
+        assert len(self.new_interaction_centers) == len(self.new_interaction_zoom_out_factors)
+        if len(self.new_interaction_centers) > 1:
+            print('It seems like more than one interaction was added since the last prediction. This is not '
+                  'recommended and may cause unexpected behavior or inefficient predictions')
+
+        start_predict = time()
+        with torch.autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
+            for prediction_center, initial_zoom_out_factor in zip(self.new_interaction_centers, self.new_interaction_zoom_out_factors):
+                # Store previous prediction for comparison
+                previous_prediction = torch.clone(self.interactions[0])
+
+                if not self.do_autozoom:
+                    initial_zoom_out_factor = 1
+
+                initial_zoom_out_factor = min(initial_zoom_out_factor, 4)
+                zoom_out_factor = initial_zoom_out_factor
+                max_zoom_out_factor = initial_zoom_out_factor
+
+                start_autozoom = time()
+                while zoom_out_factor is not None and zoom_out_factor <= 4:
+                    print('Performing prediction at zoom out factor', zoom_out_factor)
+                    max_zoom_out_factor = max(max_zoom_out_factor, zoom_out_factor)
+                    
+                    # Calculate patch size and bounding box
+                    scaled_patch_size = [round(i * zoom_out_factor) for i in self.configuration_manager.patch_size]
+                    scaled_bbox = [[int(c - p // 2), int(c + p // 2 + p % 2)] for c, p in zip(prediction_center, scaled_patch_size)]
+
+                    # Crop and prepare input
+                    crop_img, pad = crop_to_valid(self.preprocessed_image, scaled_bbox)
+                    crop_img = crop_img.to(self.device, non_blocking=self.device.type == 'cuda')
+                    crop_interactions, pad_interaction = crop_to_valid(self.interactions, scaled_bbox)
+
+                    # Resize if needed
+                    if not all([i == j for i, j in zip(self.configuration_manager.patch_size, scaled_patch_size)]):
+                        crop_interactions_resampled_gpu = torch.empty((7, *self.configuration_manager.patch_size), dtype=torch.float16, device=self.device)
+                        
+                        # Handle previous segmentation and bbox channels
+                        for i in range(0, 3):
+                            if any([x for y in pad_interaction for x in y]):
+                                tmp = pad_cropped(crop_interactions[i].to(self.device, non_blocking=self.device.type == 'cuda'), pad_interaction)
+                            else:
+                                tmp = crop_interactions[i].to(self.device)
+                            crop_interactions_resampled_gpu[i] = interpolate(tmp[None, None], self.configuration_manager.patch_size, mode='area')[0][0]
+                        empty_cache(self.device)
+
+                        # Handle point and scribble channels with dilation
+                        max_pool_ks = round_to_nearest_odd(zoom_out_factor * 2 - 1)
+                        for i in range(3, 7):
+                            if any([x for y in pad_interaction for x in y]):
+                                tmp = pad_cropped(crop_interactions[i].to(self.device, non_blocking=self.device.type == 'cuda'), pad_interaction)
+                            else:
+                                tmp = crop_interactions[i].to(self.device, non_blocking=self.device.type == 'cuda')
+                            if max_pool_ks > 1:
+                                tmp = iterative_3x3_same_padding_pool3d(tmp[None, None], max_pool_ks)[0, 0]
+                            crop_interactions_resampled_gpu[i] = interpolate(tmp[None, None], self.configuration_manager.patch_size, mode='area')[0][0]
+                        del tmp
+
+                        crop_img = interpolate(pad_cropped(crop_img, pad)[None] if any([x for y in pad_interaction for x in y]) else crop_img[None], 
+                                             self.configuration_manager.patch_size, mode='trilinear')[0]
+                        crop_interactions = crop_interactions_resampled_gpu
+                        del crop_interactions_resampled_gpu
+                        empty_cache(self.device)
+                    else:
+                        crop_img = pad_cropped(crop_img, pad) if any([x for y in pad_interaction for x in y]) else crop_img
+                        crop_interactions = pad_cropped(crop_interactions.to(self.device, non_blocking=self.device.type == 'cuda'), pad_interaction) if any([x for y in pad_interaction for x in y]) else crop_interactions.to(self.device, non_blocking=self.device.type == 'cuda')
+
+                    # Make prediction
+                    input_for_predict = torch.cat((crop_img, crop_interactions))
+                    del crop_img, crop_interactions
+                    pred = self.network(input_for_predict[None])[0].argmax(0).detach()
+                    del input_for_predict
+
+                    # Check for changes at borders
+                    previous_zoom_prediction = crop_and_pad_nd(self.interactions[0], scaled_bbox).to(self.device, non_blocking=self.device.type == 'cuda')
+                    if not all([i == j for i, j in zip(pred.shape, previous_zoom_prediction.shape)]):
+                        previous_zoom_prediction = interpolate(previous_zoom_prediction[None, None].to(float), pred.shape, mode='nearest')[0, 0]
+
+                    # Determine if we need to continue zooming
+                    continue_zoom = False
+                    if zoom_out_factor < 4 and self.do_autozoom:
+                        for dim in range(len(scaled_bbox)):
+                            if continue_zoom:
+                                break
+                            for idx in [0, pred.shape[dim] - 1]:
+                                slice_prev = previous_zoom_prediction.index_select(dim, torch.tensor(idx, device=self.device))
+                                slice_curr = pred.index_select(dim, torch.tensor(idx, device=self.device))
+                                pixels_prev = torch.sum(slice_prev)
+                                pixels_current = torch.sum(slice_curr)
+                                pixels_diff = torch.sum(slice_prev != slice_curr)
+                                rel_change = max(pixels_prev, pixels_current) / max(min(pixels_prev, pixels_current), 1e-5) - 1
+                                
+                                if pixels_diff > 1500 or (pixels_diff > 100 and rel_change > 0.2):
+                                    continue_zoom = True
+                                    if self.verbose:
+                                        print(f'Continuing zoom due to significant changes at borders')
+                                    break
+                                del slice_prev, slice_curr, pixels_prev, pixels_current, pixels_diff
+                        del previous_zoom_prediction
+
+                    # Resize prediction if needed
+                    if not all([i == j for i, j in zip(pred.shape, scaled_patch_size)]):
+                        pred = (interpolate(pred[None, None].to(float), scaled_patch_size, mode='trilinear')[0, 0] >= 0.5).to(torch.uint8)
+
+                    # Update interactions and target buffer
+                    if zoom_out_factor == 1 or not continue_zoom:
+                        pred = pred.cpu()
+                        paste_tensor(self.interactions[0], pred.half(), scaled_bbox)
+                        
+                        # Update target buffer
+                        bbox = [[i[0] + bbc[0], i[1] + bbc[0]] for i, bbc in zip(scaled_bbox, self.preprocessed_props['bbox_used_for_cropping'])]
+                        paste_tensor(self.target_buffer, pred, bbox)
+
+                    del pred
+                    empty_cache(self.device)
+
+                    if continue_zoom:
+                        zoom_out_factor *= 1.5
+                        zoom_out_factor = min(4, zoom_out_factor)
+                    else:
+                        zoom_out_factor = None
+
+                end = time()
+                print(f'Auto zoom stage took {round(end - start_autozoom, ndigits=3)}s. Max zoom out factor was {max_zoom_out_factor}')
+
+        print(f'Done. Total time {round(time() - start_predict, 3)}s')
+
+        self.new_interaction_centers = []
+        self.new_interaction_zoom_out_factors = []
+        empty_cache(self.device)
+
+    def _iterative_adjust_prediction(self, pred_prob: torch.Tensor, crop_interactions: torch.Tensor,
+                                   max_iterations: int = 15, prob_increase_factor: float = 1.5) -> torch.Tensor:
+        """
+        Perform iterative prediction adjustment when positive interactions exist.
+        
+        Args:
+            pred_prob: Probability prediction tensor [C, H, W, D]
+            crop_interactions: Interaction tensor [7, H, W, D]
+            max_iterations: Maximum number of iterations to try
+            prob_increase_factor: Factor to increase foreground probability by in each iteration
+            
+        Returns:
+            Adjusted prediction tensor
+        """
+        # Check if there are any positive interactions
+        crop_interactions_pos = crop_interactions[1:7:2]
+        pos_mask = torch.any(crop_interactions_pos > 0, dim=0)
+        pos_mask_np = pos_mask.cpu().numpy()
+        max_iterations = max_iterations if np.any(pos_mask_np) else 1
+                
+        iteration = 0
+        while iteration < max_iterations:
+            pred_prob = self._adjust_prediction_with_interactions(pred_prob, crop_interactions)
+            pred_np = pred_prob[1].cpu().numpy()
+            
+            # If prediction is all zero, try again with adjusted probabilities
+            if not np.any(pred_np):
+                # Increase foreground probability for regions with positive interactions
+                pred_prob[1, pos_mask] = torch.clamp(pred_prob[1, pos_mask] * prob_increase_factor, 0, 1)
+                pred_prob[0, pos_mask] = 1 - pred_prob[1, pos_mask]
+                iteration += 1
+            else:
+                break
+                
+        return pred_prob
+
+    def _adjust_prediction_with_interactions(self, pred_prob: torch.Tensor, crop_interactions: torch.Tensor) -> torch.Tensor:
+        """
+        Adjust prediction based on interaction masks using superpixel segmentation.
+        
+        Args:
+            pred_prob: Probability prediction tensor [C, H, W, D]
+            crop_interactions: Interaction tensor [7, H, W, D]
+            
+        Returns:
+            Adjusted prediction tensor
+        """
+        # Separate positive and negative interactions
+        crop_interactions_pos = crop_interactions[1:7:2]
+        crop_interactions_neg = crop_interactions[2:7:2]
+        
+        pos_mask = torch.any(crop_interactions_pos > 0, dim=0)
+        neg_mask = torch.any(crop_interactions_neg > 0, dim=0)
+        
+        # Separate connected components
+        import scipy.ndimage
+        from skimage.segmentation import slic
+        # Get initial prediction for labeling using threshold
+        pred_np = (pred_prob[1].cpu().numpy() > 0.5).astype(np.uint8)
+        labeled_pred, num_components = scipy.ndimage.label(pred_np)
+        
+        # Convert masks to numpy for overlap checking
+        pos_mask_np = pos_mask.cpu().numpy()
+        neg_mask_np = neg_mask.cpu().numpy()
+        
+        # Check overlap for each component and adjust pred_prob
+        for comp_id in range(1, num_components + 1):
+            comp_mask = (labeled_pred == comp_id).astype(np.uint8)
+            
+            # Check overlap with positive and negative masks
+            overlap_pos = np.logical_and(comp_mask, pos_mask_np)
+            overlap_neg = np.logical_and(comp_mask, neg_mask_np)
+            
+            # If component overlaps with both positive and negative masks
+            if np.any(overlap_pos) and np.any(overlap_neg):
+                # Get the bounding box of the component
+                bbox = scipy.ndimage.find_objects(comp_mask)[0]
+                comp_region = comp_mask[bbox]
+                pos_region = overlap_pos[bbox]
+                neg_region = overlap_neg[bbox]
+                
+                # Get pred_prob values for the region
+                pred_region_prob = pred_prob[:, bbox[0], bbox[1], bbox[2]].cpu().numpy()
+                
+                # Create RGB image from probabilities
+                pred_rgb = np.transpose(pred_region_prob, (1, 2, 3, 0))  # [H, W, D, C]
+                # pred_rgb = np.mean(pred_rgb, axis=-1, keepdims=True)  # Average across channels
+                # pred_rgb = np.repeat(pred_rgb, 3, axis=-1)  # Repeat for RGB
+                
+                # Create superpixels based on pred_prob values
+                n_segments = min(100, np.sum(comp_region))  # Limit number of segments
+                segments = slic(pred_rgb, n_segments=n_segments, compactness=10, channel_axis=-1)
+                
+                # Process each superpixel
+                for seg_id in range(1, segments.max() + 1):
+                    seg_mask = (segments == seg_id)
+                    seg_pos = np.logical_and(seg_mask, pos_region)
+                    seg_neg = np.logical_and(seg_mask, neg_region)
+                    
+                    # Get global coordinates for this segment
+                    seg_coords = np.where(seg_mask)
+                    global_coords = tuple(c + b for c, b in zip(seg_coords, [b.start for b in bbox]))
+                    
+                    # Assign values based on interaction overlap
+                    if np.any(seg_pos) and not np.any(seg_neg):
+                        pred_prob[0, global_coords] = 0.0
+                        pred_prob[1, global_coords] = 1.0
+                    elif np.any(seg_neg) and not np.any(seg_pos):
+                        pred_prob[0, global_coords] = 1.0
+                        pred_prob[1, global_coords] = 0.0
+                    # If segment has both interactions, use the original prediction
+                    else:
+                        continue
+            
+            # If component only overlaps with positive mask, force it to foreground
+            elif np.any(overlap_pos):
+                pred_prob[0, comp_mask > 0] = 0.0  # Set background to 0
+                pred_prob[1, comp_mask > 0] = 1.0  # Set foreground to 1
+            
+            # If component only overlaps with negative mask, force it to background
+            elif np.any(overlap_neg):
+                pred_prob[0, comp_mask > 0] = 1.0  # Set background to 1
+                pred_prob[1, comp_mask > 0] = 0.0  # Set foreground to 0
+                
+            # # If component does not overlap with any masks, force it to background
+            # else:
+            #     pred_prob[0, comp_mask > 0] = 1.0  # Set background to 1
+            #     pred_prob[1, comp_mask > 0] = 0.0  # Set foreground to 0
+
+        # Return thresholded prediction
+        return pred_prob
+
+
+def transform_coordinates_noresampling(
+        coords_orig: Union[List[int], Tuple[int, ...]],
+        nnunet_preprocessing_crop_bbox: List[Tuple[int, int]]
+) -> Tuple[int, ...]:
+    """
+    converts coordinates in the original uncropped image to the internal cropped representation. Man I really hate
+    nnU-Net's crop to nonzero!
+    """
+    return tuple([coords_orig[d] - nnunet_preprocessing_crop_bbox[d][0] for d in range(len(coords_orig))])
+
+