nnInteractive 2.0.0__py3-none-any.whl

nnInteractive/supervoxel/src/supervoxel.py

@@ -0,0 +1,198 @@
+import os
+import torch
+from typing import List
+import numpy as np
+from reader import NfitiReaderWriter, BloscReaderWriter
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+from sam2.sam2.build_sam import build_sam2_video_predictor
+from tqdm import tqdm
+import concurrent.futures
+import cc3d
+
+
+class SuperVoxelGenerator:
+    def __init__(self, input_dir, output_dir: str, config: dict):
+        """
+        SuperVoxelGenerator class constructor. This class is responsible for generating supervoxels
+        segmentation masks from a list of image paths.
+
+        Parameters:
+            input_dir (str): The folder containing the images to process.
+            output_dir (str): The directory where the segmentation masks will be saved. output_dir / name_of_image
+            config (dict): A dictionary containing the configuration parameters for the SuperVoxel generation.
+        """
+        self.input_dir = input_dir
+        self.output_dir = output_dir
+        self.file_format = config["file_format"]
+
+        self.reader_writer = {".nii.gz": NfitiReaderWriter(), ".b2nd": BloscReaderWriter()}[self.file_format]
+        print(f"Using {self.reader_writer} to read and write files")
+
+        # Get the list of files to process
+        self.list_of_files = [
+            f
+            for f in os.listdir(input_dir)
+            if f.endswith(config["file_format"]) and config["excluded_strings"] not in f
+        ]
+        self.config = config
+
+        self.sam = sam_model_registry["vit_h"](checkpoint=self.config["sam1_checkpoint"]).to("cuda")
+        self.mask_generator = SamAutomaticMaskGenerator(
+            model=self.sam,
+            points_per_side=48,
+            points_per_batch=256,
+            pred_iou_thresh=0.85,
+            stability_score_thresh=0.92,
+            box_nms_thresh=0.6,
+            crop_nms_thresh=0.6,
+            crop_n_layers=1,
+            crop_n_points_downscale_factor=2,
+            min_mask_region_area=192,
+        )
+        model_cfg = "sam2/configs/sam2.1/sam2.1_hiera_l.yaml"
+        self.sam2_predictor = build_sam2_video_predictor(model_cfg, ckpt_path=self.config["sam2_checkpoint"])
+
+    def sam2_propagation(self, sam2_predictor_state, image_data: np.ndarray, masks: List[np.ndarray], slice_idx: int):
+        """
+        Propagate the masks using SAM2
+
+        Parameters:
+            image_data (np.ndarray): The image data to process. Shape: (z, y, x)
+            masks (List[np.ndarray]): A list of masks to propagate. Shape: (z, y, x)
+            slice_idx (int): The index of the slice which contains the masks
+        """
+        propagated_masks = np.zeros((len(masks), *image_data.shape))
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+
+            for obj_id, m in enumerate(masks):
+                # add new prompts and instantly get the output on the same frame
+                frame_idx, object_ids, _ = self.sam2_predictor.add_new_mask(
+                    sam2_predictor_state, slice_idx, obj_id, m.astype(np.int8)
+                )
+
+            # propagate the prompts to get masklets throughout the video
+            for frame_idx, object_ids, out_masks in self.sam2_predictor.propagate_in_video(sam2_predictor_state):
+                for obj_n, out_mask in zip(object_ids, out_masks):
+                    propagated_masks[obj_n, frame_idx] = (out_mask > 0).detach().cpu().numpy()[0].astype(np.uint8)
+
+            # reset state to predict in other direction
+            self.sam2_predictor.reset_state(sam2_predictor_state)
+
+            # flip images order
+            sam2_predictor_state["images"] = torch.flip(sam2_predictor_state["images"], dims=(0,))
+            max_frame = sam2_predictor_state["images"].shape[0] - 1
+
+            for obj_id, m in enumerate(masks):
+                # add new prompts and instantly get the output on the same frame
+                frame_idx, object_ids, _ = self.sam2_predictor.add_new_mask(
+                    sam2_predictor_state, max_frame - slice_idx, obj_id, m.astype(np.int8)
+                )
+
+            # propagate the prompts to get masklets throughout the video
+            for frame_idx, object_ids, out_masks in self.sam2_predictor.propagate_in_video(sam2_predictor_state):
+                for obj_n, out_mask in zip(object_ids, out_masks):
+                    propagated_masks[obj_n, max_frame - frame_idx] = (
+                        (out_mask > 0).detach().cpu().numpy()[0].astype(np.uint8)
+                    )
+
+        return propagated_masks
+
+    def remove_other_components(self, masks, frame_idx):
+        """
+        Remove components other than the one overlaping with the seed mask in the given frame
+
+        Parameters:
+            mask (np.ndarray): The mask to process. Shape: (n, z, y, x)
+            frame_idx (int): The index of the frame
+        """
+        filtered_masks = []
+        for i, m in enumerate(masks):
+            m_cc, num_cc = cc3d.connected_components(m, connectivity=26, binary_image=True, return_N=True)
+
+            # keep all the components that overlap with the seed mask
+            filtered_components = []
+            for n in range(1, num_cc + 1):
+                if np.sum(m_cc[frame_idx] == n) > 0:
+                    filtered_components.append(n)
+            if len(filtered_components) != 0:
+                final_component = np.random.choice(filtered_components)
+                filtered_masks.append((m_cc == final_component).astype(np.uint8))
+        filtered_masks = np.stack(filtered_masks, axis=0)
+        return filtered_masks
+
+    def sam_supervoxel(self, image_data: np.ndarray):
+        """
+        Generate the supervoxels segmentation masks using SAM and SAM2
+
+        Parameters:
+            image_data (np.ndarray): The image data to process. Shape: (z, y, x)
+        """
+        # Normalize between 0 to 1 using 95 percentile
+        image_data = (image_data - np.percentile(image_data, 5)) / (
+            np.percentile(image_data, 95) - np.percentile(image_data, 5)
+        )
+        data_shape = image_data.shape
+
+        # Sample random slice with Gaussian probability
+        z_len = data_shape[0]
+        slice_probabilitys = np.exp(-np.linspace(-1, 1, data_shape[0]) ** 2 * 2)
+        slice_idx = np.random.choice(z_len, p=slice_probabilitys / slice_probabilitys.sum())
+
+        def generate_sam_masks():
+            return self.mask_generator.generate(image_data[slice_idx, ..., None].repeat(3, axis=2))
+
+        def init_sam2_predictor():
+            return self.sam2_predictor.init_state(image_data)
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future_sam_masks = executor.submit(generate_sam_masks)
+            future_sam2_predictor = executor.submit(init_sam2_predictor)
+
+            masks = future_sam_masks.result()
+            sam2_predictor_state = future_sam2_predictor.result()
+
+        # Pick n masks randomly
+        selected_masks = (
+            np.random.choice(masks, size=self.config["masks_per_image"], replace=False)
+            if len(masks) > self.config["masks_per_image"]
+            else masks
+        )
+        selected_masks = [m["segmentation"] for m in selected_masks]
+
+        # Propagate the masks using SAM2
+        propagated_masks = self.sam2_propagation(sam2_predictor_state, image_data, selected_masks, slice_idx)
+        self.sam2_predictor.reset_state(sam2_predictor_state)
+
+        # Remove other components
+        propagated_masks = self.remove_other_components(propagated_masks, slice_idx)
+
+        # Binrize the masks
+        propagated_masks = (propagated_masks > 0).astype(np.uint8)
+
+        # viewer = napari.Viewer()
+        # viewer.add_image(image_data)
+        # viewer.add_labels(propagated_masks.astype(np.uint8))
+        # napari.run()
+
+        return propagated_masks
+
+    def process_images(self):
+        """
+        Process the images in the list_of_images and generate the supervoxels segmentation masks.
+        """
+        print(f"Processing {self.list_of_files} images")
+
+        for image_name in tqdm(self.list_of_files):
+            out_path = os.path.join(self.output_dir, image_name)
+            if os.path.exists(out_path):
+                continue
+
+            image_data, metadata = self.reader_writer.read(os.path.join(self.input_dir, image_name))
+
+            if len(image_data.shape) == 4:
+                image_data = image_data[0]
+            out_img = self.sam_supervoxel(image_data)
+
+            if os.path.exists(out_path):
+                continue
+            self.reader_writer.write(out_img, metadata, out_path)

nnInteractive/trainer/nnInteractiveTrainer.py

@@ -0,0 +1,24 @@
+from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
+from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager
+from torch import nn
+
+
+class nnInteractiveTrainer_stub:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @staticmethod
+    def build_network_architecture(
+        plans_manager: PlansManager,
+        configuration_manager: ConfigurationManager,
+        num_input_channels: int,
+        num_output_channels: int,
+        enable_deep_supervision: bool = True,
+    ) -> nn.Module:
+        return nnUNetTrainer.build_network_architecture(
+            plans_manager,
+            configuration_manager,
+            num_input_channels,
+            2,  # nnunet handles one class segmentation still as CE so we need 2 outputs.
+            enable_deep_supervision,
+        )

nnInteractive/utils/bboxes.py

@@ -0,0 +1,217 @@
+from time import time
+from time import time
+from typing import List, Union, Tuple
+
+import numpy as np
+import torch
+
+
+def generate_bounding_boxes(
+    mask,
+    bbox_size=(192, 192, 192),
+    stride: Union[List[int], Tuple[int, int, int], str] = (16, 16, 16),
+    margin=(10, 10, 10),
+    max_depth=5,
+    current_depth=0,
+):
+    """
+    Generate overlapping bounding boxes to cover a 3D binary segmentation mask using PyTorch tensors.
+
+    Parameters:
+    - mask: 3D PyTorch tensor with values 0 or 1 (binary mask)
+    - bbox_size: Tuple or list of three integers specifying the size of bounding boxes per dimension (x, y, z)
+    - stride: Tuple or list of three integers specifying the stride for subsampling centers per dimension
+    - margin: Tuple or list of three integers specifying the margin to leave uncovered per dimension
+    - max_depth: Maximum recursion depth to prevent infinite recursion
+    - current_depth: Current recursion depth (used internally)
+
+    Returns:
+    - List of tuples [(min_coords, max_coords), ...], where min_coords and max_coords are lists [x, y, z] defining each box
+      as a half-open interval [min_coords, max_coords).
+    """
+    if not torch.any(mask):
+        return []
+
+    # Prevent infinite recursion
+    if current_depth > max_depth:
+        # print('random fallback due to max recursion depth')
+        return random_sampling_fallback(mask, bbox_size, margin, 25)
+
+    # Ensure bbox_size, stride, and margin are lists
+    bbox_size = list(bbox_size)
+    margin = list(margin)
+
+    # Compute half sizes for each dimension
+    half_size = [bs // 2 for bs in bbox_size]
+    # Adjust end offsets to ensure full bbox_size (handles odd sizes)
+    end_offset = [bs - hs for bs, hs in zip(bbox_size, half_size)]  # e.g., 193 - 96 = 97
+
+    # Step 1: Find all object voxels
+    object_voxels = torch.nonzero(mask, as_tuple=False)
+    if object_voxels.numel() == 0:
+        return []
+
+    # Step 2: Compute the object's bounding box to limit potential centers
+    min_coords = object_voxels.min(dim=0)[0]
+    max_coords = object_voxels.max(dim=0)[0]
+
+    if isinstance(stride, str) and stride == "auto":
+        stride = [max(1, round((j.item() - i.item()) / 4)) for i, j in zip(min_coords, max_coords)]
+
+    stride = list(stride)
+    # print('stride', stride)
+    # print('bbox', [[i, j] for i, j in zip(min_coords, max_coords)])
+
+    # Step 3: Generate potential centers within the object's bounding box
+    potential_centers = []
+    for x in range(max(0, min_coords[0].item()), min(mask.shape[0], max_coords[0].item() + 1), stride[0]):
+        for y in range(max(0, min_coords[1].item()), min(mask.shape[1], max_coords[1].item() + 1), stride[1]):
+            for z in range(max(0, min_coords[2].item()), min(mask.shape[2], max_coords[2].item() + 1), stride[2]):
+                if mask[x, y, z]:
+                    potential_centers.append([x, y, z])
+    # print(f'got {len(potential_centers)} center candidates')
+
+    if len(potential_centers) == 0:
+        return generate_bounding_boxes(
+            mask, bbox_size, [max(1, s // 2) for s in stride], margin, max_depth, current_depth + 1
+        )
+
+    potential_centers = torch.tensor(potential_centers, device=mask.device)
+
+    # Step 4: Greedy set cover algorithm
+    uncovered = mask.clone().byte()  # Use byte tensor for efficiency
+    bboxes = []
+
+    while len(potential_centers) > 0 and uncovered.any():
+        best_center = None
+        best_covered = 0
+        best_bounds = None
+
+        # Find the center that covers the most uncovered voxels
+        idx = 0
+        while idx < len(potential_centers):
+            center = potential_centers[idx]
+            c_x, c_y, c_z = center
+            x_start = max(0, c_x - half_size[0] + margin[0])
+            x_end = min(mask.shape[0], c_x + end_offset[0] - margin[0])  # Use end_offset for odd sizes
+            y_start = max(0, c_y - half_size[1] + margin[1])
+            y_end = min(mask.shape[1], c_y + end_offset[1] - margin[1])
+            z_start = max(0, c_z - half_size[2] + margin[2])
+            z_end = min(mask.shape[2], c_z + end_offset[2] - margin[2])
+
+            num_covered = uncovered[x_start:x_end, y_start:y_end, z_start:z_end].sum().item()
+            if num_covered > best_covered:
+                best_covered = num_covered
+                best_center = idx
+                best_bounds = (x_start, x_end, y_start, y_end, z_start, z_end)
+            idx += 1
+
+        # If no new voxels are covered, stop
+        if best_covered == 0:
+            break
+
+        # Add the best bounding box
+        c_x, c_y, c_z = [i.item() for i in potential_centers[best_center]]
+        bboxes.append(
+            [
+                [c_x - half_size[0], c_x + end_offset[0]],
+                [c_y - half_size[1], c_y + end_offset[1]],
+                [c_z - half_size[2], c_z + end_offset[2]],
+            ]
+        )
+
+        # Mark voxels as covered, respecting the margin
+        x_s, x_e, y_s, y_e, z_s, z_e = best_bounds
+        uncovered[
+            x_s:x_e,
+            y_s:y_e,
+            z_s:z_e,
+        ] = 0
+
+        # Remove the used center from potential_centers
+        potential_centers = potential_centers[uncovered[tuple(potential_centers.T)] > 0]
+
+    # Step 5: Recursively cover remaining voxels using uncovered as the mask
+    if uncovered.any():
+        if True:  # uncovered.sum() < np.prod([i // 3 for i in bbox_size]):
+            # print('random fallback')
+            bboxes.extend(random_sampling_fallback(uncovered, bbox_size, margin, 10))
+        else:
+            remaining_bboxes = generate_bounding_boxes(
+                uncovered, bbox_size, [max(1, s // 2) for s in stride], margin, max_depth, current_depth + 1
+            )
+            bboxes.extend(remaining_bboxes)
+
+    return bboxes
+
+
+def random_sampling_fallback(mask: torch.Tensor, bbox_size=(192, 192, 192), margin=(10, 10, 10), n_samples: int = 25):
+    half_size = [bs // 2 for bs in bbox_size]
+    # Adjust end offsets to ensure full bbox_size (handles odd sizes)
+    end_offset = [bs - hs for bs, hs in zip(bbox_size, half_size)]  # e.g., 193 - 96 = 97
+
+    bboxes = []
+
+    while mask.any():
+        indices = torch.nonzero(mask)  # nx3
+
+        best_center = None
+        best_covered = 0
+        best_bounds = None
+
+        # Find the center that covers the most uncovered voxels
+        for i in range(n_samples):
+            idx = np.random.choice(len(indices))
+            center = indices[idx]
+            c_x, c_y, c_z = [int(i.item()) for i in center]
+            x_start = max(0, c_x - half_size[0] + margin[0])
+            x_end = min(mask.shape[0], c_x + end_offset[0] - margin[0])  # Use end_offset for odd sizes
+            y_start = max(0, c_y - half_size[1] + margin[1])
+            y_end = min(mask.shape[1], c_y + end_offset[1] - margin[1])
+            z_start = max(0, c_z - half_size[2] + margin[2])
+            z_end = min(mask.shape[2], c_z + end_offset[2] - margin[2])
+
+            num_covered = mask[x_start:x_end, y_start:y_end, z_start:z_end].sum().item()
+            if num_covered > best_covered:
+                best_covered = num_covered
+                best_center = center
+                best_bounds = (x_start, x_end, y_start, y_end, z_start, z_end)
+
+        # Add the best bounding box
+        c_x, c_y, c_z = [int(i.item()) for i in best_center]
+        bboxes.append(
+            [
+                [c_x - half_size[0], c_x + end_offset[0]],
+                [c_y - half_size[1], c_y + end_offset[1]],
+                [c_z - half_size[2], c_z + end_offset[2]],
+            ]
+        )
+
+        # Mark voxels as covered, respecting the margin
+        x_s, x_e, y_s, y_e, z_s, z_e = best_bounds
+        mask[
+            x_s:x_e,
+            y_s:y_e,
+            z_s:z_e,
+        ] = 0
+    return bboxes
+
+
+if __name__ == "__main__":
+    times = []
+    torch.set_num_threads(8)
+    for _ in range(1):
+        st = time()
+        mask = torch.zeros((256, 256, 256), dtype=torch.uint8, device=0)
+        mask[50:150, 50:150, 50:150] = 1  # A cubic object
+
+        # Generate bounding boxes with an odd size to test
+        bboxes = random_sampling_fallback(
+            mask, bbox_size=(193, 193, 193), stride="auto", margin=(10, 10, 10)  # Odd size
+        )
+
+        # Print results
+        print(f"Number of bounding boxes: {len(bboxes)}")
+        end = time()
+        times.append(end - st)
+    print(times)

nnInteractive/utils/checkpoint_cleansing.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def cleanse_checkpoint_for_release(checkpoint: str):
+    a = torch.load(checkpoint, weights_only=False)
+    del a["optimizer_state"]
+    del a["init_args"]["dataset_json"]
+    a["trainer_name"] = "nnInteractiveTrainer_stub"
+    torch.save(a, checkpoint)