Merge pull request #66 from neuroneural/mindgrab

Mindgrab: powerful omnimodal skull stripping model

Author: sergeyplis

brainchop-mainthread.js

@@ -4,7 +4,9 @@ import {
   addZeroPaddingTo3dTensor,
   applyMriThreshold,
   binarizeVolumeDataTensor,
-  convByOutputChannelAndInputSlicing,
+    convByOutputChannelAndInputSlicing,
+    gn_convByOutputChannelAndInputSlicing,
+    LayerNormInPlace,
   draw3dObjBoundingVolume,
   firstLastNonZero3D,
   generateBrainMask,
@@ -191,15 +193,31 @@ async function inferenceFullVolumeSeqCovLayerPhase2(
         if (res.layers[i].activation.getClassName() !== 'linear') {
           curTensor[i] = await res.layers[i].apply(curTensor[i - 1])
         } else {
-          curTensor[i] = await convByOutputChannelAndInputSlicing(
-            curTensor[i - 1],
-            res.layers[i].getWeights()[0],
-            res.layers[i].getWeights()[1],
-            res.layers[i].strides,
-            res.layers[i].padding,
-            res.layers[i].dilationRate,
-            3
-          ) // important for memory use
+  // Check if the layer's name ends with our special suffix
+  if (res.layers[i].name.endsWith('_gn')) {
+    // Use the new GroupNorm-aware convolution function
+    curTensor[i] = await gn_convByOutputChannelAndInputSlicing(
+      curTensor[i - 1],
+      res.layers[i].getWeights()[0],
+      res.layers[i].getWeights()[1], // Can be undefined, the function handles it
+      res.layers[i].strides,
+      res.layers[i].padding,
+      res.layers[i].dilationRate,
+      3
+    );
+  } else {
+    // Use the original convolution function for non-GN layers
+    curTensor[i] = await convByOutputChannelAndInputSlicing(
+      curTensor[i - 1],
+      res.layers[i].getWeights()[0],
+      res.layers[i].getWeights()[1],
+      res.layers[i].strides,
+      res.layers[i].padding,
+      res.layers[i].dilationRate,
+      3
+    );
+  }
+            
         }
         tf.dispose(curTensor[i - 1])
       } catch (err) {
@@ -555,8 +573,13 @@ async function inferenceFullVolumePhase2(
     const curTensor = []
     curTensor[0] = cropped_slices_3d_w_pad.reshape(adjusted_input_shape)
     const timer = window.setInterval(async function () {
-      try {
-        curTensor[i] = res.layers[i].apply(curTensor[i - 1])
+        try {
+            let resultTensor = await res.layers[i].apply(curTensor[i - 1]);
+            if (res.layers[i].name.endsWith('_gn')) {
+                // LayerNormInPlace will dispose of the old resultTensor internally.
+                resultTensor = LayerNormInPlace(resultTensor);
+            }
+            curTensor[i] = resultTensor;            
       } catch (err) {
         callbackUI(err.message, -1, err.message)
         window.clearInterval(timer)
@@ -573,6 +596,7 @@ async function inferenceFullVolumePhase2(
 
         return 0
       }
+        
       callbackUI('Layer ' + i.toString(), (i + 1) / layersLength)
       console.log('layer output Tensor shape : ', curTensor[i].shape)
       console.log('layer count params ', res.layers[i].countParams())

brainchop-parameters.js

@@ -376,5 +376,51 @@ const inferenceModelsList = [
     inferenceDelay: 100, // Delay in ms time while looping layers applying.
     description:
       'FreeSurfer aparc+aseg atlas 104 parcellate brain areas into 104 regions. It contains a combination of the Desikan-Killiany atlas for cortical area and also segmentation of subcortical regions. The model use sequential convolution for inference to overcome browser memory limitations but leads to longer computation time. '
-  }
+  },
+  {
+    id: 16,
+    type: 'Brain_Extraction',
+    path: '/models/mindgrab/model.json',
+    modelName: '\u{1F9E0}\u{1FA93} omnimodal Skull Strip (High Mem, Fast)',
+    preModelId: null, // Model run first e.g.  crop the brain  { null, 1, 2, ..  }
+    preModelPostProcess: false, // If true, perform postprocessing to remove noisy regions after preModel inference generate output.
+    isBatchOverlapEnable: false, // create extra overlap batches for inference
+    numOverlapBatches: 0, // Number of extra overlap batches for inference
+    enableTranspose: true, // Keras and tfjs input orientation may need a tranposing step to be matched
+    enableCrop: true, // For speed-up inference, crop brain from background before feeding to inference model to lower memory use.
+    cropPadding: 0, // Padding size add to cropped brain
+    autoThreshold: 0, // Threshold between 0 and 1, given no preModel and tensor is normalized either min-max or by quantiles. Will remove noisy voxels around brain
+    enableQuantileNorm: false, // Some models needs Quantile Normaliztion.
+    filterOutWithPreMask: false, // Can be used to multiply final output with premodel output mask to crean noisy areas
+    enableSeqConv: false, // For low memory system and low configuration, enable sequential convolution instead of last layer
+    textureSize: 0, // Requested Texture size for the model, if unknown can be 0.
+    warning:
+      "This model may need dedicated graphics card.  For more info please check with Browser Resources <i class='fa fa-cogs'></i>.",
+    inferenceDelay: 100, // Delay in ms time while looping layers applying.
+    description:
+      'Extract the brain high accuracy model operates on full T1 image in a single pass, but uses only 11 filters per layer. Can work on dedicated graphics cards. Still more accurate than the fast version.'
+  },
+  {
+    id: 17,
+    type: 'Brain_Extraction',
+    path: '/models/mindgrab/model.json',
+    modelName: '\u{1F9E0}\u{1FA93} omnimodal Skull Strip (Low Mem, Slow)',
+    preModelId: null, // Model run first e.g.  crop the brain  { null, 1, 2, ..  }
+    preModelPostProcess: false, // If true, perform postprocessing to remove noisy regions after preModel inference generate output.
+    isBatchOverlapEnable: false, // create extra overlap batches for inference
+    numOverlapBatches: 0, // Number of extra overlap batches for inference
+    enableTranspose: true, // Keras and tfjs input orientation may need a tranposing step to be matched
+    enableCrop: true, // For speed-up inference, crop brain from background before feeding to inference model to lower memory use.
+    cropPadding: 0, // Padding size add to cropped brain
+    autoThreshold: 0, // Threshold between 0 and 1, given no preModel and tensor is normalized either min-max or by quantiles. Will remove noisy voxels around brain
+    enableQuantileNorm: false, // Some models needs Quantile Normaliztion.
+    filterOutWithPreMask: false, // Can be used to multiply final output with premodel output mask to crean noisy areas
+    enableSeqConv: true, // For low memory system and low configuration, enable sequential convolution instead of last layer
+    textureSize: 0, // Requested Texture size for the model, if unknown can be 0.
+    warning:
+      "This model may need dedicated graphics card.  For more info please check with Browser Resources <i class='fa fa-cogs'></i>.",
+    inferenceDelay: 100, // Delay in ms time while looping layers applying.
+    description:
+      'Extract the brain high accuracy model operates on  image in a single pass, but uses only 11 filters per layer. Can work on dedicated graphics cards. Still more accurate than the fast version.'
+  },    
 ] // inferenceModelsList

brainchop-webworker.js

@@ -1,24 +1,26 @@
 import * as tf from '@tensorflow/tfjs'
 import { inferenceModelsList } from './brainchop-parameters.js'
 import {
-  addZeroPaddingTo3dTensor,
-  applyMriThreshold,
-  binarizeVolumeDataTensor,
-  convByOutputChannelAndInputSlicing,
-  draw3dObjBoundingVolume,
-  firstLastNonZero3D,
-  generateBrainMask,
-  generateOutputSlicesV2,
-  getAllSlicesDataAsTF3D,
-  getModelNumLayers,
-  getModelNumParameters,
-  isModelChnlLast,
-  load_model,
-  minMaxNormalizeVolumeData,
-  quantileNormalizeVolumeData,
-  removeZeroPaddingFrom3dTensor,
-  resizeWithZeroPadding,
-  SequentialConvLayer
+    addZeroPaddingTo3dTensor,
+    applyMriThreshold,
+    binarizeVolumeDataTensor,
+    convByOutputChannelAndInputSlicing,
+    gn_convByOutputChannelAndInputSlicing,
+    LayerNormInPlace,
+    draw3dObjBoundingVolume,
+    firstLastNonZero3D,
+    generateBrainMask,
+    generateOutputSlicesV2,
+    getAllSlicesDataAsTF3D,
+    getModelNumLayers,
+    getModelNumParameters,
+    isModelChnlLast,
+    load_model,
+    minMaxNormalizeVolumeData,
+    quantileNormalizeVolumeData,
+    removeZeroPaddingFrom3dTensor,
+    resizeWithZeroPadding,
+    SequentialConvLayer
 } from './tensor-utils.js'
 
 function callbackUI(message = '', progressFrac = -1, modalMessage = '', statData = []) {
@@ -209,15 +211,31 @@ async function inferenceFullVolumeSeqCovLayerPhase2(
         if (res.layers[i].activation.getClassName() !== 'linear') {
           curTensor[i] = await res.layers[i].apply(curTensor[i - 1])
         } else {
-          curTensor[i] = await convByOutputChannelAndInputSlicing(
-            curTensor[i - 1],
-            res.layers[i].getWeights()[0],
-            res.layers[i].getWeights()[1],
-            res.layers[i].strides,
-            res.layers[i].padding,
-            res.layers[i].dilationRate,
-            3
-          ) // important for memory use
+  // Check if the layer's name ends with our special suffix
+  if (res.layers[i].name.endsWith('_gn')) {
+    // Use the new GroupNorm-aware convolution function
+    curTensor[i] = await gn_convByOutputChannelAndInputSlicing(
+      curTensor[i - 1],
+      res.layers[i].getWeights()[0],
+      res.layers[i].getWeights()[1], // Can be undefined, the function handles it
+      res.layers[i].strides,
+      res.layers[i].padding,
+      res.layers[i].dilationRate,
+      3
+    );
+  } else {
+    // Use the original convolution function for non-GN layers
+    curTensor[i] = await convByOutputChannelAndInputSlicing(
+      curTensor[i - 1],
+      res.layers[i].getWeights()[0],
+      res.layers[i].getWeights()[1],
+      res.layers[i].strides,
+      res.layers[i].padding,
+      res.layers[i].dilationRate,
+      3
+    );
+  }
+                                    
         }
 
         tf.dispose(curTensor[i - 1])
@@ -575,8 +593,12 @@ async function inferenceFullVolumePhase2(
 
     while (true) {
       try {
-        // -- curTensor[i] = res.layers[i].apply( curTensor[i-1])
-        curTensor[i] = res.layers[i].apply(curTensor[i - 1])
+          let resultTensor = await res.layers[i].apply(curTensor[i - 1]);
+          if (res.layers[i].name.endsWith('_gn')) {
+              // LayerNormInPlace will dispose of the old resultTensor internally.
+              resultTensor = LayerNormInPlace(resultTensor);
+          }
+          curTensor[i] = resultTensor;
       } catch (err) {
         callbackUI(err.message, -1, err.message)
         tf.engine().endScope()
@@ -592,10 +614,20 @@ async function inferenceFullVolumePhase2(
 
         return 0
       }
+        if (res.layers[i].activation.getClassName() == 'linear') {
+            // --- FORCE EXECUTION BY READING THE FIRST ELEMENT (Corrected) ---
+            // 1. Create a tiny slice tensor.
+            const firstElement = curTensor[i].slice([0, 0, 0, 0, 0], [1, 1, 1, 1, 1]);           
+            // 2. Awaiting its data forces GPU synchronization.
+            await firstElement.data();
+            // 3. Manually dispose of the temporary slice tensor now that it has served its purpose.
+            firstElement.dispose();
+            // --- SYNCHRONIZATION IS NOW COMPLETE ---        
+        }
       callbackUI('Layer ' + i.toString(), (i + 1) / layersLength)
       console.log('layer output Tensor shape : ', curTensor[i].shape)
       console.log('layer count params ', res.layers[i].countParams())
-      res.layers[i].dispose()
+      //res.layers[i].dispose()
       curTensor[i - 1].dispose()
       if (tf.memory().unreliable) {
         const unreliableReasons = 'unreliable reasons :' + tf.memory().reasons

public/models/mindgrab/colormap.json

@@ -0,0 +1,6 @@
+{
+    "R": [ 0, 255],
+    "G": [ 0, 255],
+    "B": [ 0, 255],
+    "labels": [ "background", "brain mask"]
+}