melage 0.0.68__py3-none-any.whl → 1.0.0__py3-none-any.whl

melage/plugins/warpseg/warpseg_main/data_reader/utils.py

@@ -0,0 +1,267 @@
+import numpy as np
+from skimage.measure import label as label_connector
+import nibabel as nib
+import SimpleITK as sitk
+
+def createPathData_IXI():
+    return
+def creaePathData_LDM_fullhead():
+    return
+
+code_direction = (('L', 'R'), ('P', 'A'), ('I', 'S'))
+
+def read_atlas_file(atlas_file, printing=False, save_again=False):
+    im = nib.load(atlas_file)
+    if im.ndim == 4 and im.shape[-1]==1:
+        im = nib.funcs.four_to_three(im)[0]
+    orig_orient = nib.io_orientation(im.affine)
+    code_direction = (('L', 'R'), ('P', 'A'), ('I', 'S'))
+    source_system = ''.join(list(aff2axcodes(im.affine, code_direction)))
+    if source_system != 'RAS':
+        print(f'converted to RAS {atlas_file}')
+        target_orient = axcodes2ornt('RAS', code_direction)
+        transform = ornt_transform(orig_orient, target_orient)
+        im = im.as_reoriented(transform)
+        if save_again:
+            im.to_filename(atlas_file)
+
+    if printing:
+        print('{} with range {:.2f}, {:.2f}'.format(atlas_file, im.get_fdata().min(), im.get_fdata().max()))
+    #atlas_ims[..., r] = im.get_fdata().copy()
+    return im
+
+
+def read_sitk_as_nib(sitk_im):
+    return nib.Nifti1Image(sitk.GetArrayFromImage(sitk_im).transpose(),
+                         make_affine(sitk_im), None)
+
+
+def make_affine(simpleITKImage):
+    # https://niftynet.readthedocs.io/en/v0.2.1/_modules/niftynet/io/simple_itk_as_nibabel.html
+    # get affine transform in LPS
+    if simpleITKImage.GetDimension() == 4:
+        c = [simpleITKImage.TransformContinuousIndexToPhysicalPoint(p)
+             for p in ((1, 0, 0, 0),
+                       (0, 1, 0, 0),
+                       (0, 0, 1, 0),
+                       (0, 0, 0, 0))]
+        c = np.array(c)
+        c = c[:, :-1]
+    elif simpleITKImage.GetDimension() == 3:
+        c = [simpleITKImage.TransformContinuousIndexToPhysicalPoint(p)
+             for p in ((1, 0, 0),
+                       (0, 1, 0),
+                       (0, 0, 1),
+                       (0, 0, 0))]
+        c = np.array(c)
+    affine = np.concatenate([
+        np.concatenate([c[0:3] - c[3:], c[3:]], axis=0),
+        [[0.], [0.], [0.], [1.]]
+    ], axis=1)
+    affine = np.transpose(affine)
+    # convert to RAS to match nibabel
+    affine = np.matmul(np.diag([-1., -1., 1., 1.]), affine)
+    return affine
+
+def resample_to_size(im, new_size=None, method='linear',scale_factor=None):
+    original_image = read_nib_as_sitk(im)
+    # Get the current size of the image
+    size = original_image.GetSize()
+    spacing = original_image.GetSpacing()
+    # Calculate the scale factor for resizing
+    if scale_factor is not None:
+        # Compute new size if scale_factor is provided
+        new_size = [int(round(orig_sz * sf)) for orig_sz, sf in zip(size, scale_factor)]
+        scale_factor = [orig_sp / sf for orig_sp, sf in zip(spacing, scale_factor)]
+    else:
+        scale_factor = [(float(sz)/new_sz)*spc for sz, new_sz, spc in zip(size, new_size, original_image.GetSpacing())]
+
+    # Resample the image using the scale factor
+    resampler = sitk.ResampleImageFilter()
+    resampler.SetReferenceImage(original_image)
+    resampler.SetOutputSpacing(scale_factor)
+    resampler.SetSize([int(el) for el in new_size])
+
+    resampler.SetOutputOrigin(original_image.GetOrigin())
+    resampler.SetOutputDirection(original_image.GetDirection())
+
+    if method.lower() == 'linear':
+        resampler.SetInterpolator(sitk.sitkLinear)  # You can choose different interpolators
+    elif method.lower() == 'nearest':
+        resampler.SetInterpolator(sitk.sitkNearestNeighbor)
+    else:
+        resampler.SetInterpolator(sitk.sitkBSpline)  # You can choose different interpolators
+
+    # Perform resampling
+    resized_image = resampler.Execute(original_image)
+    #inverse_scale_factor = [s1/float(sz) for sz, s1 in zip(*[resized_image.GetSize(), size])]
+    return read_sitk_as_nib(resized_image)#, inverse_scale_factor, list(size)
+
+def remove_zero(f_data, value=0, min_max=None):
+    """
+    Remove non segmented areas from image
+    :param f_data:
+    :param value:
+    :return:
+    """
+
+
+    shape_original = f_data.shape
+    if min_max is None:
+        min_max = []
+        xs, ys, zs = np.where(f_data > value)  # find zero values
+        tol = 4
+        for x in [xs, ys, zs]:
+            minx = min(x)-tol if min(x)-tol>1 else min(x)
+            maxx = max(x) + tol if max(x) + tol < f_data.shape[0]-1 else max(x)
+            min_max.append([minx, maxx])
+    f_data = f_data[min_max[0][0]:min_max[0][1] + 1, min_max[1][0]:min_max[1][1] + 1, min_max[2][0]:min_max[2][1] + 1]
+
+    return f_data, min_max, shape_original
+
+
+def restore_zero(cropped_data, original_shape, min_max, value=0):
+    """
+    Restore the cropped data back to the original shape with padding (default = 0).
+    :param cropped_data: The cropped array from remove_zero
+    :param original_shape: The shape of the original full image
+    :param min_max: The min/max indices used for cropping
+    :param value: The padding value to fill
+    :return: Restored full-size array
+    """
+    restored = np.full(original_shape, value, dtype=cropped_data.dtype)
+
+    restored[
+        min_max[0][0]:min_max[0][1] + 1,
+        min_max[1][0]:min_max[1][1] + 1,
+        min_max[2][0]:min_max[2][1] + 1
+    ] = cropped_data
+
+    return restored
+
+def pad_to_multiple_of_32(img, target_size = [224, 256, 192], val_const=0):
+    #target_size = [224, 256, 192]  # Nearest multiples of 32
+     # Nearest multiples of 32
+    diff = [t - s for t, s in zip(target_size, img.shape)]
+
+    # Distribute padding symmetrically
+    pad_width = [(diff[0] // 2, diff[0] - diff[0] // 2),  # Depth
+                 (diff[1] // 2, diff[1] - diff[1] // 2),  # Height
+                 (diff[2] // 2, diff[2] - diff[2] // 2)]  # Width
+
+    return np.pad(img, pad_width, mode='constant', constant_values=val_const), pad_width
+
+
+def scalecrop(data, dst_min, dst_max, src_min, scale):
+    """
+    Function to crop the intensity ranges to specific min and max values
+    :param np.ndarray data: Image data (intensity values)
+    :param float dst_min: future minimal intensity value
+    :param float dst_max: future maximal intensity value
+    :param float src_min: minimal value to consider from source (crops below)
+    :param float scale: scale value by which source will be shifted
+    :return: np.ndarray data_new: scaled image data
+    """
+    data_new = dst_min + scale * (data - src_min)
+
+    # clip
+    data_new = np.clip(data_new, dst_min, dst_max)
+    #print("Output:   min: " + format(data_new.min()) + "  max: " + format(data_new.max()))
+
+    return data_new
+
+def getscale(data, dst_min, dst_max, f_low=0.0, f_high=0.999):
+    """
+    Function to get offset and scale of image intensities to robustly rescale to range dst_min..dst_max.
+    Equivalent to how mri_convert conforms images.
+    :param np.ndarray data: image data (intensity values)
+    :param float dst_min: future minimal intensity value
+    :param float dst_max: future maximal intensity value
+    :param f_low: robust cropping at low end (0.0 no cropping)
+    :param f_high: robust cropping at higher end (0.999 crop one thousandths of high intensity voxels)
+    :return: float src_min: (adjusted) offset
+    :return: float scale: scale factor
+    """
+    # get min and max from source
+    src_min = np.min(data)
+    src_max = np.max(data)
+
+    #print("Input:    min: " + format(src_min) + "  max: " + format(src_max))
+
+    if f_low == 0.0 and f_high == 1.0:
+        return src_min, 1.0
+
+    # compute non-zeros and total vox num
+    nz = (np.abs(data) >= 1e-15).sum()
+    voxnum = data.shape[0] * data.shape[1] * data.shape[2]
+
+    # compute histogram
+    histosize = 1000
+    bin_size = (src_max - src_min) / histosize
+    hist, bin_edges = np.histogram(data, histosize)
+
+    # compute cummulative sum
+    cs = np.concatenate(([0], np.cumsum(hist)))
+
+    # get lower limit
+    nth = int(f_low * voxnum)
+    idx = np.where(cs < nth)
+
+    if len(idx[0]) > 0:
+        idx = idx[0][-1] + 1
+
+    else:
+        idx = 0
+
+    src_min = idx * bin_size + src_min
+
+    # print("bin min: "+format(idx)+"  nth: "+format(nth)+"  passed: "+format(cs[idx])+"\n")
+    # get upper limit
+    nth = voxnum - int((1.0 - f_high) * nz)
+    idx = np.where(cs >= nth)
+
+    if len(idx[0]) > 0:
+        idx = idx[0][0] - 2
+
+    else:
+        idx = 0
+        print('ERROR: rescale upper bound not found')
+
+    src_max = idx * bin_size + src_min
+    # print("bin max: "+format(idx)+"  nth: "+format(nth)+"  passed: "+format(voxnum-cs[idx])+"\n")
+
+    # scale
+    if src_min == src_max:
+        scale = 1.0
+
+    else:
+        scale = (dst_max - dst_min) / (src_max - src_min)
+
+    #print("rescale:  min: " + format(src_min) + "  max: " + format(src_max) + "  scale: " + format(scale))
+
+    return src_min, scale
+
+def normalize(img, mask=None):
+    if mask is None:
+        src_min, scale = getscale(img, 0, 255)
+        new_data = scalecrop(img, 0, 255, src_min, scale)
+
+    else:
+        src_min, scale = getscale(img*mask, 0, 255)
+        new_data = np.zeros_like(img)
+        new_data[mask] = scalecrop(img[mask], 0, 255, src_min, scale)
+        new_data[~mask] = scalecrop(img[~mask], 0, 255, src_min, scale)
+    new_data = new_data / 255
+
+
+    return new_data
+
+def LargestCC(segmentation, connectivity=3):
+    """
+    Get largest connected components
+    """
+    if segmentation.ndim == 4:
+        segmentation = segmentation.squeeze(-1)
+    labels = label_connector(segmentation, connectivity=connectivity)
+    frequency = np.bincount(labels.flat)
+    return labels, frequency

melage/plugins/warpseg/warpseg_main/dist_utils.py

@@ -0,0 +1,18 @@
+import os
+import torch.distributed as dist
+import torch
+def setup_dist(rank, world_size):
+    os.environ['MASTER_ADDR'] = 'localhost'  # Set this to the address of the master node
+    os.environ['MASTER_PORT'] = '12353'       # Set this to an available port
+    #dist.init_process_group(backend='gloo', rank=rank, world_size=world_size)
+    dist.init_process_group(
+        backend="nccl",          # Use 'nccl' for GPUs, 'gloo' for CPU
+        init_method="env://",    # URL specifying how to initialize (e.g., environment variables)
+        world_size=world_size,   # Total number of processes
+        rank=rank                # Rank of the current process
+    )
+    torch.cuda.set_device(rank)  # Set the current GPU device
+
+
+def cleanup():
+    dist.destroy_process_group()

melage/plugins/warpseg/warpseg_main/requirements.txt

@@ -0,0 +1,10 @@
+torch
+numpy
+scikit-image
+monai
+matplotlib
+scikit-learn
+scipy
+nibabel
+voxelmorph
+pkbar

melage/plugins/warpseg/warpseg_main/test.py

@@ -0,0 +1,272 @@
+import os
+import sys
+
+# Add project root to Python path to resolve module imports
+#project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
+#if project_root not in sys.path:
+#    sys.path.insert(0, project_root)
+
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import nibabel as nib
+from .voxelmorph.torch.networks import Multi_Segmentor
+from .data_reader.utils import normalize as normalize_image, read_atlas_file
+
+# Set environment variables
+os.environ['VXM_BACKEND'] = 'pytorch'
+
+
+from scipy.spatial import cKDTree
+import scipy.ndimage as ndi
+from .data_reader.utils import LargestCC
+
+def post_process(seg_whole, ind_whole):
+    """
+    Post-process segmentation to remove small disconnected components 
+    and fill holes using nearest neighbor interpolation.
+    """
+    whole_segs = []
+    initial_segs = []
+    label = seg_whole.copy()
+    
+    # 1. Filter connected components
+    for c, el in enumerate(ind_whole):
+        cur_im = (seg_whole.squeeze() == el)
+        initial_segs.append(cur_im)
+        
+        im_l, cc = LargestCC(cur_im, 1)
+        sorted_indices = [idx for idx in np.argsort(cc)[::-1] if idx != 0]
+        
+        class_mask = np.zeros_like(cur_im, dtype=bool)
+        
+        if len(sorted_indices) > 0:
+            threshold = cc[sorted_indices[0]] * 0.05 
+            selected_sorted = [idx for idx in sorted_indices if cc[idx] > threshold]
+            
+            temp_mask = np.zeros_like(im_l, dtype=int)
+            for sl in selected_sorted:
+                temp_mask += (im_l == sl).astype(int)
+            
+            class_mask = temp_mask > 0
+            
+            class_mask = ndi.binary_fill_holes(class_mask)
+        else:
+            class_mask = cur_im
+            
+        whole_segs.append(class_mask)
+
+    initial_segs = np.stack(initial_segs).astype(int)
+    whole_segs = np.stack(whole_segs).astype(int)
+    
+    # 2. Identify changed pixels (removed regions)
+    ind_change_whole = (initial_segs - whole_segs) > 0
+    
+    # 3. Fill removed regions with nearest valid labels
+    # Get all valid pixels (those not in the regions being processed or kept regions)
+    # Actually, we want to fill using labels from the *final* valid segmentation.
+    # But `label` currently has the original values.
+    # We should zero out the removed regions in `label` first? 
+    # The original code used `label` (original) and found nearest from `valid_indices`.
+    # `valid_indices` were defined as pixels in `label` NOT in `ind_whole`.
+    # This means it fills using ONLY labels that were NOT part of the post-processed set.
+    # This seems to be the intent: fill "holes" (removed bad components) with surrounding "background" or other tissues.
+    
+    label_indices = np.argwhere(label > 0)
+    label_values = label[tuple(label_indices.T)]
+    mask_valid = ~np.isin(label_values, ind_whole)
+    valid_indices = label_indices[mask_valid]
+    valid_values = label_values[mask_valid]
+    
+    if len(valid_indices) > 0:
+        tree = cKDTree(valid_indices)
+        
+        for ind in range(len(ind_whole)):
+            # Pixels that were removed for this class
+            indices = np.argwhere(ind_change_whole[ind])
+            if len(indices) > 0:
+                dist, nearest_idx = tree.query(indices)
+                nearest_labels = valid_values[nearest_idx]
+                label[tuple(indices.T)] = nearest_labels
+                
+    return label
+
+def load_model(model, optimizer, device, load_filepath):
+    """
+    Loads model weights from a file.
+    """
+    try:
+        print(f"Loading model from: {load_filepath}")
+        data = torch.load(load_filepath, map_location=device)
+    except FileNotFoundError:
+        print(f"Error: Model file not found at {load_filepath}")
+        return 0
+
+    # Determine the key for the state dictionary
+
+    pretrained_state = data["model"]
+
+
+
+    # Using strict=False is flexible for mismatched keys
+    incompatible_keys = model.load_state_dict(pretrained_state, strict=False)
+    if incompatible_keys.missing_keys:
+        print(f"Warning: Missing keys in state_dict: {incompatible_keys.missing_keys}")
+    if incompatible_keys.unexpected_keys:
+        print(f"Warning: Unexpected keys in state_dict: {incompatible_keys.unexpected_keys}")
+    #except Exception as e:
+    #    print(f"Error while loading state dict: {e}")
+        #return 0
+    
+    current_step = 0
+    try:
+        if optimizer is not None:
+            optimizer.load_state_dict(data['opt'])
+        current_step = data['step']
+    except KeyError:
+        print("Optimizer state or step not found in checkpoint. Starting from step 0.")
+    return 1
+
+def save_segmentation(segmentation, affine, base_name, out_folder, suffix='_seg_ours.nii.gz'):
+    """
+    Saves the segmentation result as a NIfTI file.
+    """
+    if '.nii.gz' in base_name:
+        new_filename = base_name.replace('.nii.gz', suffix)
+    elif '.nii' in base_name:
+        new_filename = base_name.replace('.nii', suffix)
+    else:
+        new_filename = base_name + suffix
+        
+    out_file_path = os.path.join(out_folder, new_filename)
+    nifti_image = nib.Nifti1Image(segmentation, affine, header=None)
+    nifti_image.to_filename(out_file_path)
+    print(f"Saved: {out_file_path}")
+
+def get_inference(model, input_data, device, inshape=[224, 256, 192],
+                  post_processing=True):
+    input_data = normalize_image(input_data)
+
+    # Prepare tensor
+    input_tensor = torch.from_numpy(input_data).float().to(device)
+    # Add batch and channel dimensions: (D, H, W) -> (1, 1, D, H, W)
+    # Assuming input is 3D volume
+    if input_tensor.ndim == 3:
+        input_tensor = input_tensor.unsqueeze(0).unsqueeze(0)
+    elif input_tensor.ndim == 4:  # Maybe already has channel?
+        input_tensor = input_tensor.unsqueeze(0)
+
+    source_shape = input_tensor.shape[2:]
+
+    # Resize to model input shape
+    input_tensor_resized = F.interpolate(input_tensor, size=inshape, mode='trilinear', align_corners=False)
+
+    print(f"Running inference...")
+    with torch.inference_mode():
+        # Inference
+        pred_logits, pred_logits_tissue = model(input_tensor_resized)
+
+        # Resize back to original shape
+        pred_logits = F.interpolate(pred_logits, size=source_shape, mode='trilinear', align_corners=False)
+        pred_logits_tissue = F.interpolate(pred_logits_tissue, size=source_shape, mode='trilinear', align_corners=False)
+
+        # Get probabilities and segmentations
+        prob_value = F.softmax(pred_logits, dim=1)
+        prob_value_tissue = F.softmax(pred_logits_tissue, dim=1)
+
+        segmentation = prob_value.argmax(1).squeeze().numpy().astype(np.int16)
+        if post_processing:
+            segmentation = post_process(segmentation, ind_whole=[3, 19])
+        segmentation_tissue = prob_value_tissue.argmax(1).squeeze().numpy().astype(np.int16)
+        if post_processing:
+            segmentation_tissue = post_process(segmentation_tissue, ind_whole=[3])
+
+        # Convert to numpy
+        #out_segmentation = segmentation.detach().squeeze().cpu().numpy().astype(np.int16)
+        #out_segmentation_tissue = segmentation_tissue.detach().squeeze().cpu().numpy().astype(np.int16)
+    return segmentation_tissue, segmentation
+
+
+def build_model(model_path, device):
+    inshape = [224, 256, 192]
+    model = Multi_Segmentor(tissue_roi=6, nrois=30).to(device)
+
+    # Load model
+    if model_path is not None:
+        step = load_model(model, None, device, load_filepath=model_path)
+        if step == 0:
+             raise RuntimeError(f"Failed to load model weights from {model_path}")
+    model.to(device)
+    model.eval()
+    return model
+
+def get_options():
+    parser = argparse.ArgumentParser()
+
+    # Input/Output parameters
+    parser.add_argument('--input-image', required=True, help='Path to the input image for inference')
+    parser.add_argument('--output-dir', default='test_results', help='Directory to save results')
+    parser.add_argument('--model-path', required=True, help='Path to the trained model checkpoint')
+
+    # Model parameters (must match training config)
+    parser.add_argument('--enc', type=int, nargs='+', help='list of unet encoder filters (default: 16 32 32 32)')
+    parser.add_argument('--dec', type=int, nargs='+',
+                        help='list of unet decorder filters (default: 32 32 32 32 32 16 16)')
+    parser.add_argument('--int-steps', type=int, default=7, help='number of integration steps (default: 7)')
+    parser.add_argument('--int-downsize', type=int, default=2,
+                        help='flow downsample factor for integration (default: 2)')
+    parser.add_argument('--bidir', action='store_true', help='enable bidirectional cost function')
+
+    # Other parameters
+    parser.add_argument('--gpu', default='0', help='GPU ID number(s), comma-separated (default: 0)')
+    parser.add_argument('--modelName', type=str, default='t1_fuzzy',
+                        help='Automatic')  # Kept for compatibility if needed
+
+    opt = parser.parse_args()
+    return opt
+
+def main():
+    import voxelmorph as vxm
+    opt = get_options()
+    
+    # Device configuration
+    if torch.cuda.is_available():
+        device = 'cuda:0'
+    else:
+        device = 'cpu'
+
+    # Model configuration
+    inshape = [224, 256, 192]
+    model = build_model(opt.model_path, device)
+
+    # Prepare output directory
+    os.makedirs(opt.output_dir, exist_ok=True)
+
+    # Load and process input image
+    if not os.path.exists(opt.input_image):
+        print(f"Error: Input image not found at {opt.input_image}")
+        return
+
+    try:
+        # Load image
+        nifti_img = read_atlas_file(opt.input_image)
+        input_data = nifti_img.get_fdata()
+        affine = nifti_img.affine
+        
+        # Normalize
+        out_segmentation_tissue, out_segmentation = get_inference(model, input_data, device, inshape)
+            
+        # Save results
+        base_name = os.path.basename(opt.input_image)
+        save_segmentation(out_segmentation, affine, base_name, opt.output_dir, suffix='_seg_ours.nii.gz')
+        save_segmentation(out_segmentation_tissue, affine, base_name, opt.output_dir, suffix='_tissue_seg_ours.nii.gz')
+            
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    main()