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

melage/plugins/mga_net/MGA_Net.py

@@ -0,0 +1,145 @@
+from PyQt5.QtWidgets import QMessageBox
+from melage.widgets import MelagePlugin
+from melage.dialogs.dynamic_gui import DynamicDialog
+from .MGA_Net_schema import get_schema
+from PyQt5 import QtWidgets
+from .main.test_mgaNet import build_model, get_inference
+import torch
+from PyQt5.QtCore import Qt
+from melage.config import settings
+import os
+from PyQt5.QtCore import pyqtSignal
+# --- THE LOGIC CLASS ---
+class MGA_NetLogic(DynamicDialog):
+    """
+    This class handles the BRAIN of the WarpSeg tool.
+    The LOOKS are handled automatically by DynamicDialog + Schema.
+    """
+    completed = pyqtSignal(object)
+    def __init__(self, data_context,parent=None):
+        # 1. Initialize DynamicDialog with the Schema
+        # This single line builds the entire window!
+
+        super().__init__(parent)
+        self.create_main_ui(schema=get_schema())
+        self.setAttribute(Qt.WA_DeleteOnClose)
+        self.data_context = data_context
+
+        # 2. AUTOMATICALLY WIDGET BINDING
+        # Iterate over all widgets created by the schema and bind them to 'self'.
+        for widget_id, widget_obj in self.widgets.items():
+            setattr(self, widget_id, widget_obj)
+
+        # 3. GENERAL PURPOSE SIGNAL CONNECTION (Auto-Connect)
+        # Instead of a hardcoded dictionary, we iterate through every widget from the schema
+        # and look for a matching method in this class named: on_<ID>_<Signal>
+        # This works for ANY schema changes automatically.
+        supported_signals = [
+            "clicked", "toggled", "currentIndexChanged",
+            "textChanged", "valueChanged"
+        ]
+
+        for widget_id, widget_obj in self.widgets.items():
+            for signal_name in supported_signals:
+                # 1. Check if the widget has this signal (e.g., Button has 'clicked')
+                if hasattr(widget_obj, signal_name):
+                    # 2. Check if WE have a handler method (e.g., 'on_btn_apply_clicked')
+                    handler_name = f"on_{widget_id}_{signal_name}"
+                    if hasattr(self, handler_name):
+                        # 3. Connect them!
+                        signal = getattr(widget_obj, signal_name)
+                        handler = getattr(self, handler_name)
+                        # Disconnect first to be safe (idempotent), then connect
+                        try:
+                            signal.disconnect(handler)
+                        except TypeError:
+                            pass
+                        signal.connect(handler)
+                        # print(f"Auto-connected: {widget_id}.{signal_name} -> {handler_name}")
+
+
+    @property
+    def ui_schema(self):
+        # We call the function to get the dictionary
+        return get_schema()
+
+
+
+
+    # Renamed from 'run_process' to match the schema ID 'btn_apply'
+    def on_btn_apply_clicked(self):
+        view = self.combo_view.currentText()
+        data_view = self.data_context[view]
+        if  data_view is None:
+            QMessageBox.information(self, "Error", "No image data available for the selected view.")
+            return
+
+        try:
+            """The main execution function."""
+            # Get values using auto-bound attributes
+            use_cuda = self.check_cuda.isChecked()
+
+
+            self.progress_bar.setValue(10)
+            print(f"Running WarpSeg: View={view}, CUDA={use_cuda}")
+
+            if use_cuda and torch.cuda.is_available():
+                device = torch.device("cuda")
+                print("Using CUDA for computation.")
+            else:
+                device = torch.device("cpu")
+                print("Using CPU for computation.")
+            #if self.check_adult.isChecked():
+            model_path = getattr(self, 'weights_path_custom', None)
+            #else:
+            #    model_path = getattr(self, 'weights_path_custom', None) or os.path.join(settings.DEFAULT_MODELS_DIR, "WarpSeg_Infant.pth")
+            self.progress_bar.setValue(30)
+            # Build model
+            model = build_model(model_path=model_path, device=device)
+            self.progress_bar.setValue(50)
+            # Here you would load pre-trained weights if available
+            # For simplicity, we skip that step
+            # Run inference
+
+            data = data_view.get_fdata().copy()
+            if len(data.shape)<=1:
+                raise ValueError("Input image must be at least 2D NIfTI.")
+            threshold = self.spin_threshold.value()
+            eco_mri = -1
+            if self.radio_mri.isChecked():
+                eco_mri = 1
+            image, seg = get_inference(model, data, device, eco_mri=eco_mri, threshold=threshold, high_quality_rec=True)
+            result_package = {
+                "image": image,
+                "label": seg,
+                "view": view
+            }
+            self.completed.emit(result_package)
+            self.progress_bar.setValue(100)
+            #QMessageBox.information(self, "Done", "Segmentation Complete")
+
+        except Exception as e:
+            QMessageBox.information(self, "Error", f"{e}")
+            self.progress_bar.setValue(0)
+
+    def on_context_action(self, text, widget):
+        """Handle the Right-Click Context Menu defined in schema"""
+        if text == "Reset Adult Options":
+            # Access radio button directly by ID
+            self.radio_adult_whole.setChecked(True)
+            print("Options reset.")
+
+
+# --- THE PLUGIN WRAPPER ---
+class WarpSegPlugin(MelagePlugin):
+    @property
+    def name(self) -> str: return "MGA-Net"
+
+    @property
+    def category(self) -> str: return "Segmentation"
+
+
+    def get_widget(self, data_context =None,parent=None):
+        logic = MGA_NetLogic(data_context, parent)
+
+        return logic

melage/plugins/mga_net/MGA_Net_schema.py

@@ -0,0 +1,29 @@
+
+from ..ui_helpers import Label, Combo, Check, Radio, Group, Button, Progress, HBox, FilePicker, SpinBox, Reference
+
+# The simplified schema using helper functions.
+# The main layout defaults to 'vbox', so items stack automatically.
+def get_schema():
+    UI_SCHEMA = {
+        "title": "MGA-Net Infant Deep Learning Segmentation",
+        "min_width": 500,
+        "layout": "vbox",
+        "items": [
+            #Label(id="lbl_model_info", text="Model Info: Ready for Ifant brain segmentation"),
+            Reference(
+                "<b>Jafrasteh et al. (2024)</b>. <i>'A novel mask-guided attention...'</i>. "
+                '<a href="https://www.sciencedirect.com/science/article/pii/S1053811924003690">(paper)</a>'
+
+            ),
+            SpinBox(id="spin_threshold", label="Segmentation Threshold:", value=0.5, min_val=-4.0, max_val=4.0, step=0.1, decimals=1),
+            # Context menu allows for extra actions like resetting options
+
+            Group(id="group_mode", layout="hbox", children=[
+                Radio(id="radio_mri", text="MRI Segmentation", checked=True),
+                Radio(id="radio_us", text="US Segmentation")
+            ]),
+
+
+        ]
+    }
+    return UI_SCHEMA

melage/plugins/mga_net/main/figures/Network.txt

@@ -0,0 +1 @@
+

melage/{widgets/DeepLModels/new_unet.py → plugins/mga_net/main/model/mga_net.py}

@@ -637,7 +637,7 @@ class UnetGen(nn.Module):
 
         return [mask,z]
 
-class Unet3D(UnetGen):
+class MGA_NET(UnetGen):
     def __init__(self,time_embed=False, channels=1, *args, encoder_class=UnetEncoder, **kwargs):
         self.time_embed = time_embed
         self.use_tr_conv = False
@@ -647,6 +647,6 @@ class Unet3D(UnetGen):
                          block=ResidualBlock, norm_type=norm_type, **kwargs)
 
         self.channels = channels
-        self.netName = 'Unet3D'
+        self.netName = 'MGA_NET'
     def name(self):
-        return 'unet3d'
+        return 'MGA_NET'

melage/plugins/mga_net/main/model/utils.py

@@ -0,0 +1,258 @@
+import numpy as np
+import nibabel as nib
+import SimpleITK as sitk
+from skimage.measure import label as label_connector
+__AUTHOR__ = 'Bahram Jafrasteh'
+code_direction = (('L', 'R'), ('P', 'A'), ('I', 'S'))
+def convert_to_ras(affine, target = "RAS"):
+    """
+    Args:
+        affine: affine matrix
+        target: target system
+
+    Returns:
+
+    """
+    from nibabel.orientations import aff2axcodes, axcodes2ornt, ornt_transform
+    orig_orient = nib.io_orientation(affine)
+    source_system = ''.join(list(aff2axcodes(affine, code_direction)))# get direction
+    target_orient = axcodes2ornt(target, code_direction)
+    transform = ornt_transform(orig_orient, target_orient)
+
+    return transform, source_system
+def LargestCC(segmentation, connectivity=3):
+    """
+    Get largets connected components
+    """
+    if segmentation.ndim == 4:
+        segmentation = segmentation.squeeze(-1)
+    labels = label_connector(segmentation, connectivity=connectivity)
+    frequency = np.bincount(labels.flat)
+    return labels, frequency
+
+
+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 read_sitk_as_nib(sitk_im):
+    return nib.Nifti1Image(sitk.GetArrayFromImage(sitk_im).transpose(),
+                         make_affine(sitk_im), None)
+def read_nib_as_sitk(image_nib, dtype=None):
+    # From https://github.com/gift-surg/PySiTK/blob/master/pysitk/simple_itk_helper.py
+    if dtype is None:
+        dtype = np.float32#image_nib.header["bitpix"].dtype
+    nda_nib = image_nib.get_fdata().astype(dtype)
+    nda_nib_shape = nda_nib.shape
+    nda = np.zeros((nda_nib_shape[2],
+                    nda_nib_shape[1],
+                    nda_nib_shape[0]),
+                   dtype=dtype)
+
+    # Convert to (Simple)ITK data array format, i.e. reorder to
+    # z-y-x-components shape
+    for i in range(0, nda_nib_shape[2]):
+        for k in range(0, nda_nib_shape[0]):
+            nda[i, :, k] = nda_nib[k, :, i]
+    # Get SimpleITK image
+    vector_image_sitk = sitk.GetImageFromArray(nda)
+    # Update header from nibabel information
+    # (may introduce some header inaccuracies?)
+    R = np.array([
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, 0, 1]])
+    affine_nib = image_nib.affine.astype(np.float64)
+    R_nib = affine_nib[0:-1, 0:-1]
+
+    spacing_sitk = np.array(image_nib.header.get_zooms(), dtype=np.float64)
+    spacing_sitk = spacing_sitk[0:R_nib.shape[0]]
+    S_nib_inv = np.diag(1. / spacing_sitk)
+
+    direction_sitk = R.dot(R_nib).dot(S_nib_inv).flatten()
+
+    t_nib = affine_nib[0:-1, 3]
+    origin_sitk = R.dot(t_nib)
+
+    vector_image_sitk.SetSpacing(np.array(spacing_sitk).astype('double'))
+    vector_image_sitk.SetDirection(direction_sitk)
+    vector_image_sitk.SetOrigin(origin_sitk)
+    return vector_image_sitk
+def resample_to_size(im, new_size, scale_factor=None,method='linear'):
+    """
+    Resample image to new size
+    """
+    original_image = read_nib_as_sitk(im)
+    # Get the current size of the image
+    size = original_image.GetSize()
+
+    # Calculate the scale factor for resizing
+    if scale_factor is None:
+        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
+    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 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_mri(img):
+    # taken from npp
+    #https://github.com/Novestars/Neural_Pre_Processing/blob/master/nppy/models/utils.py
+    src_min, scale = getscale(img, 0, 255)
+    new_data = scalecrop(img, 0, 255, src_min, scale)
+    return new_data
+
+def remove_zero(f_data, value=0):
+    """
+    Remove non segmented areas from image
+    :param f_data:
+    :param value:
+    :return:
+    """
+
+    xs, ys, zs = np.where(f_data > value) #find zero values
+    tol = 4
+
+    min_max = []
+    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
+
+def get_back_data(im, shape_initial, pad_zero, border_value):
+    """
+    Get data back to its original shape
+    """
+    im_fill = np.ones(shape_initial) * border_value
+    im_fill[pad_zero[0][0]:pad_zero[0][1] + 1, pad_zero[1][0]:pad_zero[1][1] + 1,
+    pad_zero[2][0]:pad_zero[2][1] + 1] = im
+    return im_fill

melage/plugins/mga_net/main/test_mgaNet.py

@@ -0,0 +1,134 @@
+__AUTHOR__ = 'Bahram Jafrasteh'
+from .model.mga_net import MGA_NET
+
+import torch
+import sys
+import os
+from model.utils import *
+from scipy.ndimage import binary_fill_holes
+import torch.nn.functional as F
+
+
+
+
+
+################## Loading model ################################
+def build_model(model_path="MGA_NET.pth", device="cpu"):
+    model = MGA_NET(time_embed=True)
+    model.to(device)
+    if model_path is None:
+        model_path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "MGA_NET.pth")
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict['model'], strict=True)
+    return model.eval()
+
+
+
+################## Loading data and data standardization ################################
+def get_inference(model, imA, device, eco_mri=1, threshold=0.0, high_quality_rec=True):
+    time = torch.from_numpy(np.array(eco_mri)).unsqueeze(0).to(torch.float).to(device)
+
+    border_value = imA[0, 0, 0]
+    shape_initial = imA.shape
+    image_used, pad_zero = remove_zero(imA, border_value)
+
+    image_used = normalize_mri(image_used)/255.0
+    shape_zero = image_used.shape
+
+    target_shape = [128, 128, 128]
+    #image_used_1 = resample_to_size(nib.Nifti1Image(image_used, affine, header), new_size=target_shape,
+    #                                   method='spline').get_fdata()
+    image_used_torch = torch.from_numpy(image_used)[None, None, ...].to(torch.float).to(device)
+    imB = F.interpolate(image_used_torch, size=target_shape, mode='trilinear', align_corners=False)
+
+
+    ################## Brain extraction and image reconstruction ################################
+    im_low = model.forward(imB, time)
+    im_mask_low, im_rec_low = im_low
+
+    if high_quality_rec:
+
+        target_shape = [192, 192, 192] # to create higher quality images
+        #image_used_2 = resample_to_size(nib.Nifti1Image(image_used, affine, header), new_size=target_shape,
+        #                          method='spline').get_fdata()
+        image_used_torch = torch.from_numpy(image_used)[None, None, ...].to(torch.float).to(device)
+        im_high_input = F.interpolate(image_used_torch, size=target_shape, mode='trilinear', align_corners=False)
+
+        #imA = torch.from_numpy(image_used_2).to(torch.float).unsqueeze(0).unsqueeze(0)
+        #imA = imA.to(device)
+        im_high = model.forward(im_high_input, time)
+
+        im_mask_high, im_rec_high = im_high
+
+
+        im_mask = im_mask_low.detach().cpu()
+        im_rec = im_rec_high.detach().cpu()
+    else:
+        im_mask = im_mask_low.detach().cpu()
+        im_rec = im_rec_low.detach().cpu()
+
+    ################## resmaple to the original size ################################
+    #im_mask = resample_to_size(nib.Nifti1Image(im_mask, affine, header), new_size=shape_zero,
+    #                           method='spline').get_fdata()
+    im_mask = F.interpolate(im_mask, size=shape_zero, mode='nearest')
+    im_rec = F.interpolate(im_rec, size=shape_zero, mode='trilinear', align_corners=False)
+    #im_rec = resample_to_size(nib.Nifti1Image(im_rec, affine, header), new_size=shape_zero,
+    #                          method='spline').get_fdata()
+    im_mask = im_mask.detach().squeeze().cpu().numpy()
+    im_rec = im_rec.detach().squeeze().cpu().numpy()
+
+    im_mask = get_back_data(im_mask, shape_initial, pad_zero, im_mask[0, 0, 0])
+    im_rec = get_back_data(im_rec, shape_initial, pad_zero, im_rec[0, 0, 0])
+
+
+    ind = im_mask >= threshold
+    im_mask[ind] = 0
+    im_mask[~ind] = 1
+
+    im_mask = binary_fill_holes(im_mask)
+    im_mask, labels_freq = LargestCC(im_mask, connectivity=1)
+    if len(labels_freq)>2:
+        ind_argmax = np.argmax(
+            [imA[im_mask == el].sum() for el in range(len(labels_freq)) if el != 0]) + 1
+        ind = im_mask != ind_argmax
+        im_mask[ind] = 0
+        im_mask[~ind] = 1
+    im_rec = normalize_mri(im_rec)
+
+    return im_rec, im_mask
+
+def make_ras_image(imA):
+    transform, source = convert_to_ras(imA.affine, target='RAS')
+    if source != 'RAS':
+        imA = imA.as_reoriented(transform)
+    return imA
+def main():
+    file_inp = sys.argv[1]
+    eco_mri = int(sys.argv[2])  # -1 for US and 1 for MRI
+    threshold = 0.0
+    if len(sys.argv) > 3:
+        threshold = float(sys.argv[3])
+    high_quality_rec = True  # Network has been trained on 128x128x128 size image. However, it is possible to sample 192x192x192 images to get higher quality images
+    basen = os.path.basename(file_inp)
+    basen = basen[:basen.find('.nii')]
+
+    file_inp_mask = os.path.join(os.path.dirname(file_inp), basen + '_mask.nii.gz')
+    file_inp_rec = os.path.join(os.path.dirname(file_inp), basen + '_rec.nii.gz')
+    if torch.cuda.is_available():
+        # device = torch.device("cuda")
+        device = torch.device("cuda:0")
+        torch.cuda.set_device(0)
+        torch.backends.cudnn.benchmark = True
+    # else:
+    device = torch.device("cpu")
+    model = build_model("MGA_NET.pth", device)
+    imA = nib.load(file_inp)
+
+    imA = make_ras_image(imA)
+
+    affine = imA.affine
+    header = imA.header
+    imB, im_mask = get_inference(model, imA, device)
+
+    imB = nib.Nifti1Image(imB, affine, header).to_filename(file_inp_rec)
+    im_mask = nib.Nifti1Image(im_mask, affine, header).to_filename(file_inp_mask)