RepConTUR / preprocessing.py

Add RepConTUR-1B weights, remote code, tokenizer and model card

c6be546 verified 6 days ago

12.4 kB

"""
MONAI-based 3D preprocessing pipeline for VoxTell grounding volumes.
Pipeline (matches VoxTell's nnUNet-style training preprocessing):
  1. EnsureChannelFirstd — add channel dim
  2. Orientationd(RAS)  — standard anatomical orientation
  3. Spacingd((1,1,1))  — isotropic 1 mm voxel spacing
  4. NormalizeIntensityd — z-score normalization (mean/std over all voxels)
  5. SpatialPadd(192)   — ensure volume is at least 192^3
  6. RandCropByPosNegLabeld(192^3, pos=1, neg=1) — foreground-biased crop
Replaces MonaiSAM2Preprocessor for the VoxTell segmentation backbone.
"""

from typing import List, Optional, Tuple

import numpy as np
import torch
from ._compat import print_log

from monai.data import MetaTensor
from monai.transforms import (
    Compose,
    EnsureChannelFirstd,
    NormalizeIntensityd,
    Orientationd,
    RandCropByPosNegLabeld,
    SpatialPadd,
    CenterSpatialCropd,
    Spacingd,
    Transposed,
    SpatialPad
)
from monai.utils import ensure_tuple


class RandCropByPosNegLabelWithRandomOffsetd(RandCropByPosNegLabeld):
    """
    Random Cropping with consistent output spacing and size, using random offset.
    Pipeline: Pad -> RandomCrop (with random offset) -> Resize to target_spacing -> CenterCrop to spatial_size
    """
    
    def randomize(
        self,
        label: torch.Tensor,
        fg_indices: np.ndarray | None = None,
        bg_indices: np.ndarray | None = None,
        image: torch.Tensor | None = None,
    ) -> None:
        # Call parent's randomize to set self.centers
        super().randomize(label, fg_indices, bg_indices, image)

        if self.cropper.centers is None:
            return

        # Vectorized random offset and clamping
        centers = np.array(self.cropper.centers)  # (num_centers, num_dims)
        spatial_size = np.array(ensure_tuple(self.cropper.spatial_size))
        spatial_shape = np.array(label.shape[1:])

        # Generate random offsets for all centers at once: [0, size) - size//2
        offsets = (self.R.rand(*centers.shape) * spatial_size).astype(np.int64) - spatial_size // 2

        # Apply offsets and clamp to valid bounds
        new_centers = centers + offsets
        min_centers = spatial_size // 2
        max_centers = spatial_shape - spatial_size + spatial_size // 2
        new_centers = np.clip(new_centers, min_centers, max_centers).astype(np.int64)

        # Convert back to tuple of tuples (required format)
        self.cropper.centers = tuple(map(tuple, new_centers.tolist()))
        
        
class SmartRandCropByPosNegLabeld(RandCropByPosNegLabelWithRandomOffsetd):
    """
    Random Cropping with consistent output spacing and size and random offset.
    Pipeline: Pad -> RandomCrop -> Resize to target_spacing -> CenterCrop to spatial_size
    """
    
    def __init__(
        self,
        keys,
        label_key,
        spatial_size,
        target_spacing,
        pad_mode="constant",
        resize_mode=('bilinear', 'nearest'),
        **kwargs
    ):
        """
        Args:
            keys: keys of the corresponding items to be transformed.
            label_key: key for label used for positive/negative sampling.
            spatial_size: final output spatial size after all transforms.
            target_spacing: target spacing to resample to (e.g., [1.0, 1.0, 1.0]).
            pad_mode: padding mode for SpatialPadd.
            resize_mode: interpolation mode for Spacingd ("bilinear", "nearest", etc.).
            **kwargs: additional arguments for RandCropByPosNegLabeld.
        """
        self.spatial_size = np.array(spatial_size)
        self.target_spacing = np.array(target_spacing)
        self.pad_mode = pad_mode
        
        # Initialize parent with dummy spatial_size, will be computed per-image
        super().__init__(keys=keys, label_key=label_key, spatial_size=(1, 1, 1), **kwargs)
        
        self.spacer = Spacingd(
            keys=self.keys,
            pixdim=self.target_spacing.tolist(),
            mode=resize_mode
        )
        self.center_cropper = CenterSpatialCropd(keys=self.keys, roi_size=self.spatial_size.tolist())
    
    def _get_crop_size(self, label):
        """Compute crop size in voxels based on physical size and image spacing."""        
        if isinstance(label, MetaTensor):
            spacing = torch.norm(label.affine[:3, :3], dim=0)
        else:
            print_log(f"Image has no affine, using default spacing of 1.0", logger='current')
            spacing = self.target_spacing
        
        # Compute crop voxels: physical_extent / current_spacing
        # physical_extent = spatial_size[i] * target_spacing[i]
        crop_size = [torch.ceil(self.spatial_size[i] * (self.target_spacing[i] / spacing[i])).int().item() for i in range(3)]
        crop_size = [s + 5 for s in crop_size]
        
        return tuple(crop_size)
    
    def __call__(self, data, lazy=None):
        d = dict(data)
        label = d[self.label_key]
        
        # Step 1: Compute crop size based on image spacing
        crop_size = self._get_crop_size(label)
        self.cropper.spatial_size = crop_size # This is used for random offset sampling
        
        # Step 2: Pad if volume is smaller than crop_size
        for key in self.key_iterator(d):
            img = d[key]
            if self.pad_mode == 'min_value' or self.pad_mode == 'constant':
                pad_mode = 'constant'
                
                constant_values = img.min() if self.pad_mode == 'min_value' else 0
                
                padder = SpatialPad(spatial_size=crop_size, mode=pad_mode, constant_values=constant_values)
            else:
                pad_mode = self.pad_mode
                
                padder = SpatialPad(spatial_size=crop_size, mode=pad_mode)
            
            d[key] = padder(img)
        
        # Step 3: Random crop (uses parent's __call__) with random offset
        cropped_list = super().__call__(d, lazy=lazy)
        
        # From list of dicts to dict of lists
        stacked_list = []
        for sample in cropped_list:
            # Step 4: Resample to target spacing
            sample = self.spacer(sample)
            # Step 5: Center crop to exact spatial size
            sample = self.center_cropper(sample)
            stacked_list.append(sample)
            
        return stacked_list
            

class VoxTellPreprocessor:
    """
    3D preprocessing for VoxTell grounding input.
    Takes a raw CT volume + binary mask (any orientation, any spacing)
    and produces one or more 192^3 patches ready for VoxTell's encoder.
    In **eval mode** (set via ``set_eval_mode()``), returns the full
    preprocessed volume (orient → z-score → resample 1mm → transpose)
    without any cropping, for sliding window inference.
    Args:
        patch_size: Spatial size of each output patch.
        pos: Weight for positive (foreground) patches.
        neg: Weight for negative (background) patches.
        num_samples: Number of 192^3 patches to sample per volume.
    """

    def __init__(
        self,
        patch_size: Tuple[int, int, int] = (192, 192, 192),
        pos: float = 1.0,
        neg: float = 1.0,
        num_samples: int = 1,
    ):
        self.patch_size = patch_size
        self.num_samples = num_samples
        self._eval_mode = False

        self.transform = Compose([
            # 1. Add channel dimension: (W,H,D) → (1,W,H,D)
            EnsureChannelFirstd(
                keys=['image', 'label'],
                channel_dim='no_channel',
            ),
            # 2. Reorient to RAS (Right-Anterior-Superior)
            Orientationd(
                keys=['image', 'label'],
                axcodes='RAS',
            ),
            # 3. Z-score normalization over all voxels
            NormalizeIntensityd(
                keys=['image']
            ),
            # 4. Random crop with foreground bias
            SmartRandCropByPosNegLabeld(
                keys=['image', 'label'],
                label_key='label',
                spatial_size=patch_size,
                target_spacing=(1.0, 1.0, 1.0),
                resize_mode=('bilinear', 'nearest'),
                pos=pos,
                neg=neg,
                num_samples=num_samples,
            ),
            # 5. Reorient to RAS (Right-Anterior-Superior)
            Transposed(
                keys=['image', 'label'],
                indices=[0, 3, 2, 1] # CWHD to CDHW
            ),
        ])

        # Eval pipeline: same as training but no cropping.
        # orient → normalize → resample 1mm → transpose
        self.eval_transform = Compose([
            EnsureChannelFirstd(
                keys=['image', 'label'],
                channel_dim='no_channel',
            ),
            Orientationd(
                keys=['image', 'label'],
                axcodes='RAS',
            ),
            NormalizeIntensityd(
                keys=['image'],
            ),
            Spacingd(
                keys=['image', 'label'],
                pixdim=(1.0, 1.0, 1.0),
                mode=('bilinear', 'nearest'),
            ),
            Transposed(
                keys=['image', 'label'],
                indices=[0, 3, 2, 1],  # CWHD to CDHW
            ),
        ])

    def set_eval_mode(self):
        """Switch to eval mode: full-volume preprocessing without cropping."""
        self._eval_mode = True

    def apply_image(
        self,
        volume: np.ndarray,
        mask: np.ndarray,
        affine: Optional[np.ndarray] = None,
        data_source: Optional[str] = None,
    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """
        Preprocess volume + mask and return VoxTell-ready 3D patches.
        In eval mode, returns the full preprocessed volume (no cropping)
        for sliding window inference.  In training mode, returns one or
        more foreground-biased 192^3 patches.
        Args:
            volume: Raw CT volume, any axis order (matching affine).
            mask: Binary segmentation mask, same shape as volume.
            affine: 4x4 affine matrix (required for orientation/spacing).
            data_source: Unused (kept for interface compatibility).
        Returns:
            Tuple of ``(volume_tensor, mask_tensor)`` where shapes are:
              - Training: ``(1, D, H, W)`` with D=H=W=192.
              - Eval: ``(1, D', H', W')`` at 1 mm isotropic spacing.
            If ``num_samples > 1`` (training only), returns a list of
            such tuples.  Returns None on failure.
        """
        try:
            assert affine is not None, 'Affine matrix is required'

            img_tensor = MetaTensor(
                volume.astype(np.float32), affine=affine,
            )
            lbl_tensor = MetaTensor(
                mask.astype(np.float32), affine=affine,
            )

            data_dict = {
                'image': img_tensor,
                'label': lbl_tensor,
            }

            if self._eval_mode:
                result = self.eval_transform(data_dict)
                vol = result['image'].as_tensor()  # (1, D', H', W')
                msk = result['label'].as_tensor()
                return vol, msk

            result = self.transform(data_dict)

            # RandCropByPosNegLabeld returns a list of dicts
            if isinstance(result, list):
                if self.num_samples == 1:
                    r = result[0]
                    vol_patch = r['image'].as_tensor()  # (1, D, H, W)
                    msk_patch = r['label'].as_tensor()
                    return vol_patch, msk_patch
                else:
                    patches = []
                    for r in result:
                        vol_patch = r['image'].as_tensor()
                        msk_patch = r['label'].as_tensor()
                        patches.append((vol_patch, msk_patch))
                    return patches
            else:
                vol_patch = result['image'].as_tensor()
                msk_patch = result['label'].as_tensor()
                return vol_patch, msk_patch

        except Exception as e:
            print_log(
                f'VoxTellPreprocessor error: {e}', logger='current',
            )
            import traceback
            print_log(traceback.format_exc(), logger='current')
            return None

	"""
	MONAI-based 3D preprocessing pipeline for VoxTell grounding volumes.

	Pipeline (matches VoxTell's nnUNet-style training preprocessing):
	1. EnsureChannelFirstd — add channel dim
	2. Orientationd(RAS) — standard anatomical orientation
	3. Spacingd((1,1,1)) — isotropic 1 mm voxel spacing
	4. NormalizeIntensityd — z-score normalization (mean/std over all voxels)
	5. SpatialPadd(192) — ensure volume is at least 192^3
	6. RandCropByPosNegLabeld(192^3, pos=1, neg=1) — foreground-biased crop

	Replaces MonaiSAM2Preprocessor for the VoxTell segmentation backbone.
	"""

	from typing import List, Optional, Tuple

	import numpy as np
	import torch
	from ._compat import print_log

	from monai.data import MetaTensor
	from monai.transforms import (
	Compose,
	EnsureChannelFirstd,
	NormalizeIntensityd,
	Orientationd,
	RandCropByPosNegLabeld,
	SpatialPadd,
	CenterSpatialCropd,
	Spacingd,
	Transposed,
	SpatialPad
	)
	from monai.utils import ensure_tuple


	class RandCropByPosNegLabelWithRandomOffsetd(RandCropByPosNegLabeld):
	"""
	Random Cropping with consistent output spacing and size, using random offset.
	Pipeline: Pad -> RandomCrop (with random offset) -> Resize to target_spacing -> CenterCrop to spatial_size
	"""

	def randomize(
	self,
	label: torch.Tensor,
	fg_indices: np.ndarray \| None = None,
	bg_indices: np.ndarray \| None = None,
	image: torch.Tensor \| None = None,
	) -> None:
	# Call parent's randomize to set self.centers
	super().randomize(label, fg_indices, bg_indices, image)

	if self.cropper.centers is None:
	return

	# Vectorized random offset and clamping
	centers = np.array(self.cropper.centers) # (num_centers, num_dims)
	spatial_size = np.array(ensure_tuple(self.cropper.spatial_size))
	spatial_shape = np.array(label.shape[1:])

	# Generate random offsets for all centers at once: [0, size) - size//2
	offsets = (self.R.rand(centers.shape) spatial_size).astype(np.int64) - spatial_size // 2

	# Apply offsets and clamp to valid bounds
	new_centers = centers + offsets
	min_centers = spatial_size // 2
	max_centers = spatial_shape - spatial_size + spatial_size // 2
	new_centers = np.clip(new_centers, min_centers, max_centers).astype(np.int64)

	# Convert back to tuple of tuples (required format)
	self.cropper.centers = tuple(map(tuple, new_centers.tolist()))


	class SmartRandCropByPosNegLabeld(RandCropByPosNegLabelWithRandomOffsetd):
	"""
	Random Cropping with consistent output spacing and size and random offset.
	Pipeline: Pad -> RandomCrop -> Resize to target_spacing -> CenterCrop to spatial_size
	"""

	def __init__(
	self,
	keys,
	label_key,
	spatial_size,
	target_spacing,
	pad_mode="constant",
	resize_mode=('bilinear', 'nearest'),
	**kwargs
	):
	"""
	Args:
	keys: keys of the corresponding items to be transformed.
	label_key: key for label used for positive/negative sampling.
	spatial_size: final output spatial size after all transforms.
	target_spacing: target spacing to resample to (e.g., [1.0, 1.0, 1.0]).
	pad_mode: padding mode for SpatialPadd.
	resize_mode: interpolation mode for Spacingd ("bilinear", "nearest", etc.).
	**kwargs: additional arguments for RandCropByPosNegLabeld.
	"""
	self.spatial_size = np.array(spatial_size)
	self.target_spacing = np.array(target_spacing)
	self.pad_mode = pad_mode

	# Initialize parent with dummy spatial_size, will be computed per-image
	super().__init__(keys=keys, label_key=label_key, spatial_size=(1, 1, 1), **kwargs)

	self.spacer = Spacingd(
	keys=self.keys,
	pixdim=self.target_spacing.tolist(),
	mode=resize_mode
	)
	self.center_cropper = CenterSpatialCropd(keys=self.keys, roi_size=self.spatial_size.tolist())

	def _get_crop_size(self, label):
	"""Compute crop size in voxels based on physical size and image spacing."""
	if isinstance(label, MetaTensor):
	spacing = torch.norm(label.affine[:3, :3], dim=0)
	else:
	print_log(f"Image has no affine, using default spacing of 1.0", logger='current')
	spacing = self.target_spacing

	# Compute crop voxels: physical_extent / current_spacing
	# physical_extent = spatial_size[i] * target_spacing[i]
	crop_size = [torch.ceil(self.spatial_size[i] * (self.target_spacing[i] / spacing[i])).int().item() for i in range(3)]
	crop_size = [s + 5 for s in crop_size]

	return tuple(crop_size)

	def __call__(self, data, lazy=None):
	d = dict(data)
	label = d[self.label_key]

	# Step 1: Compute crop size based on image spacing
	crop_size = self._get_crop_size(label)
	self.cropper.spatial_size = crop_size # This is used for random offset sampling

	# Step 2: Pad if volume is smaller than crop_size
	for key in self.key_iterator(d):
	img = d[key]
	if self.pad_mode == 'min_value' or self.pad_mode == 'constant':
	pad_mode = 'constant'

	constant_values = img.min() if self.pad_mode == 'min_value' else 0

	padder = SpatialPad(spatial_size=crop_size, mode=pad_mode, constant_values=constant_values)
	else:
	pad_mode = self.pad_mode

	padder = SpatialPad(spatial_size=crop_size, mode=pad_mode)

	d[key] = padder(img)

	# Step 3: Random crop (uses parent's __call__) with random offset
	cropped_list = super().__call__(d, lazy=lazy)

	# From list of dicts to dict of lists
	stacked_list = []
	for sample in cropped_list:
	# Step 4: Resample to target spacing
	sample = self.spacer(sample)
	# Step 5: Center crop to exact spatial size
	sample = self.center_cropper(sample)
	stacked_list.append(sample)

	return stacked_list


	class VoxTellPreprocessor:
	"""
	3D preprocessing for VoxTell grounding input.

	Takes a raw CT volume + binary mask (any orientation, any spacing)
	and produces one or more 192^3 patches ready for VoxTell's encoder.

	In eval mode (set via ``set_eval_mode()``), returns the full
	preprocessed volume (orient → z-score → resample 1mm → transpose)
	without any cropping, for sliding window inference.

	Args:
	patch_size: Spatial size of each output patch.
	pos: Weight for positive (foreground) patches.
	neg: Weight for negative (background) patches.
	num_samples: Number of 192^3 patches to sample per volume.
	"""

	def __init__(
	self,
	patch_size: Tuple[int, int, int] = (192, 192, 192),
	pos: float = 1.0,
	neg: float = 1.0,
	num_samples: int = 1,
	):
	self.patch_size = patch_size
	self.num_samples = num_samples
	self._eval_mode = False

	self.transform = Compose([
	# 1. Add channel dimension: (W,H,D) → (1,W,H,D)
	EnsureChannelFirstd(
	keys=['image', 'label'],
	channel_dim='no_channel',
	),
	# 2. Reorient to RAS (Right-Anterior-Superior)
	Orientationd(
	keys=['image', 'label'],
	axcodes='RAS',
	),
	# 3. Z-score normalization over all voxels
	NormalizeIntensityd(
	keys=['image']
	),
	# 4. Random crop with foreground bias
	SmartRandCropByPosNegLabeld(
	keys=['image', 'label'],
	label_key='label',
	spatial_size=patch_size,
	target_spacing=(1.0, 1.0, 1.0),
	resize_mode=('bilinear', 'nearest'),
	pos=pos,
	neg=neg,
	num_samples=num_samples,
	),
	# 5. Reorient to RAS (Right-Anterior-Superior)
	Transposed(
	keys=['image', 'label'],
	indices=[0, 3, 2, 1] # CWHD to CDHW
	),
	])

	# Eval pipeline: same as training but no cropping.
	# orient → normalize → resample 1mm → transpose
	self.eval_transform = Compose([
	EnsureChannelFirstd(
	keys=['image', 'label'],
	channel_dim='no_channel',
	),
	Orientationd(
	keys=['image', 'label'],
	axcodes='RAS',
	),
	NormalizeIntensityd(
	keys=['image'],
	),
	Spacingd(
	keys=['image', 'label'],
	pixdim=(1.0, 1.0, 1.0),
	mode=('bilinear', 'nearest'),
	),
	Transposed(
	keys=['image', 'label'],
	indices=[0, 3, 2, 1], # CWHD to CDHW
	),
	])

	def set_eval_mode(self):
	"""Switch to eval mode: full-volume preprocessing without cropping."""
	self._eval_mode = True

	def apply_image(
	self,
	volume: np.ndarray,
	mask: np.ndarray,
	affine: Optional[np.ndarray] = None,
	data_source: Optional[str] = None,
	) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
	"""
	Preprocess volume + mask and return VoxTell-ready 3D patches.

	In eval mode, returns the full preprocessed volume (no cropping)
	for sliding window inference. In training mode, returns one or
	more foreground-biased 192^3 patches.

	Args:
	volume: Raw CT volume, any axis order (matching affine).
	mask: Binary segmentation mask, same shape as volume.
	affine: 4x4 affine matrix (required for orientation/spacing).
	data_source: Unused (kept for interface compatibility).

	Returns:
	Tuple of ``(volume_tensor, mask_tensor)`` where shapes are:
	- Training: ``(1, D, H, W)`` with D=H=W=192.
	- Eval: ``(1, D', H', W')`` at 1 mm isotropic spacing.
	If ``num_samples > 1`` (training only), returns a list of
	such tuples. Returns None on failure.
	"""
	try:
	assert affine is not None, 'Affine matrix is required'

	img_tensor = MetaTensor(
	volume.astype(np.float32), affine=affine,
	)
	lbl_tensor = MetaTensor(
	mask.astype(np.float32), affine=affine,
	)

	data_dict = {
	'image': img_tensor,
	'label': lbl_tensor,
	}

	if self._eval_mode:
	result = self.eval_transform(data_dict)
	vol = result['image'].as_tensor() # (1, D', H', W')
	msk = result['label'].as_tensor()
	return vol, msk

	result = self.transform(data_dict)

	# RandCropByPosNegLabeld returns a list of dicts
	if isinstance(result, list):
	if self.num_samples == 1:
	r = result[0]
	vol_patch = r['image'].as_tensor() # (1, D, H, W)
	msk_patch = r['label'].as_tensor()
	return vol_patch, msk_patch
	else:
	patches = []
	for r in result:
	vol_patch = r['image'].as_tensor()
	msk_patch = r['label'].as_tensor()
	patches.append((vol_patch, msk_patch))
	return patches
	else:
	vol_patch = result['image'].as_tensor()
	msk_patch = result['label'].as_tensor()
	return vol_patch, msk_patch

	except Exception as e:
	print_log(
	f'VoxTellPreprocessor error: {e}', logger='current',
	)
	import traceback
	print_log(traceback.format_exc(), logger='current')
	return None