New (hidden) feature for scalar models

Author: neurolabusc

brainchop-webworker.js

@@ -357,6 +357,10 @@ async function inferenceFullVolumeSeqCovLayerPhase2(
             opts,
             niftiImage
           )
+          if (modelEntry.isScalar) {
+            console.log(':::: This model ignores isScalar:::::::::::::::::::::::::')
+            modelEntry.isScalar = false
+          }
           console.log(' Phase-2 num of tensors after generateOutputSlicesV2: ', tf.memory().numTensors)
 
           tf.dispose(outLabelVolume)
@@ -568,7 +572,7 @@ async function inferenceFullVolumePhase2(
     statData.Model_Layers = await getModelNumLayers(res)
     statData.Model = modelEntry.modelName
     // statData.Extra_Info = null
-
+    const isScalar = modelEntry.isScalar === true
     const curTensor = []
     curTensor[0] = cropped_slices_3d_w_pad.reshape(adjusted_input_shape)
     // console.log("curTensor[0] :", curTensor[0].dataSync())
@@ -619,7 +623,21 @@ async function inferenceFullVolumePhase2(
         try {
           const argMaxTime = performance.now()
           console.log(' Try tf.argMax for fullVolume ..')
-          prediction_argmax = tf.argMax(curTensor[i], axis)
+          if (isScalar) {
+            const input = tf.softmax(curTensor[i], -1) // shape: [..., C], with C >= 2
+            const shape = input.shape
+            const lastDim = shape.length - 1
+            // Slice the last dimension to keep only channels 1 and onward (ignore channel 0)
+            const start = Array(lastDim).fill(0).concat(1)
+            const size = shape.slice(0, lastDim).concat(shape[lastDim] - 1)
+            const sliced = input.slice(start, size)
+            // Sum across the last dimension (collapsing the remaining classes)
+            const summed = tf.sum(sliced, lastDim)
+            // Remove any leading singleton dimensions (optional)
+            prediction_argmax = summed.squeeze() // only if you want to remove shape [1, ...]
+          } else {
+            prediction_argmax = tf.argMax(curTensor[i], axis)
+          }
           console.log('tf.argMax for fullVolume takes : ', ((performance.now() - argMaxTime) / 1000).toFixed(4))
         } catch (err1) {
           // if channel last
@@ -691,10 +709,12 @@ async function inferenceFullVolumePhase2(
         statData.Expect_Labels = expected_Num_labels
         statData.NumLabels_Match = numSegClasses === expected_Num_labels
 
-        if (numSegClasses !== expected_Num_labels) {
-          // errTxt = "expected " + expected_Num_labels + " labels, but the predicted are " + numSegClasses + ". For possible solutions please refer to <a href='https://github.com/neuroneural/brainchop/wiki/FAQ#Q3' target='_blank'><b> FAQ </b></a>.", "alert-error"
-          const errTxt = 'expected ' + expected_Num_labels + ' labels, but the predicted are ' + numSegClasses
-          callbackUI(errTxt, -1, errTxt)
+        if (!isScalar) {
+          if (numSegClasses !== expected_Num_labels) {
+            // errTxt = "expected " + expected_Num_labels + " labels, but the predicted are " + numSegClasses + ". For possible solutions please refer to <a href='https://github.com/neuroneural/brainchop/wiki/FAQ#Q3' target='_blank'><b> FAQ </b></a>.", "alert-error"
+            const errTxt = 'expected ' + expected_Num_labels + ' labels, but the predicted are ' + numSegClasses
+            callbackUI(errTxt, -1, errTxt)
+          }
         }
 
         // -- Transpose back to original unpadded size
@@ -722,11 +742,20 @@ async function inferenceFullVolumePhase2(
         )
         console.log(' outLabelVolume final shape after resizing :  ', outLabelVolume.shape)
 
+        if (isScalar) {
+          const thresh  = tf.scalar(0.04);   // threshold
+          const scale255 = tf.scalar(255.0); // if you still need the scaling step
+          const mask = outLabelVolume.greaterEqual(thresh).toFloat();
+          outLabelVolume = outLabelVolume.mul(mask);
+          outLabelVolume = outLabelVolume.mul(scale255);
+        }
         const filterOutWithPreMask = modelEntry.filterOutWithPreMask
         // To clean the skull area wrongly segmented in phase-2.
-        if (pipeline1_out != null && opts.isBrainCropMaskBased && filterOutWithPreMask) {
-          const bin = binarizeVolumeDataTensor(pipeline1_out)
-          outLabelVolume = outLabelVolume.mul(bin)
+        if (!isScalar) {
+          if (pipeline1_out != null && opts.isBrainCropMaskBased && filterOutWithPreMask) {
+            const bin = binarizeVolumeDataTensor(pipeline1_out)
+            outLabelVolume = outLabelVolume.mul(bin)
+          }
         }
 
         startTime = performance.now()

index.html

@@ -28,6 +28,9 @@
     <select id="modelSelect">
     </select>
     &nbsp;
+    <label for="scalarCheck" hidden>Scalar</label>
+    <input type="checkbox" id="scalarCheck" hidden unchecked />
+    &nbsp;
     <button disabled id="createMeshBtn">Create Mesh</button>
     &nbsp;
     <br>

main.js

@@ -83,6 +83,7 @@ async function main() {
     await ensureConformed()
     let model = inferenceModelsList[this.selectedIndex]
     model.isNvidia = false
+    model.isScalar = scalarCheck.checked
     const rendererInfo = nv1.gl.getExtension("WEBGL_debug_renderer_info")
     if (rendererInfo) {
       model.isNvidia = nv1.gl
@@ -151,6 +152,9 @@ async function main() {
       nv1.setClipPlane([2, 0, 90])
     }
   }
+  scalarCheck.onchange = function () {
+    modelSelect.selectedIndex = -1
+  }
   function doLoadImage() {
     saveBtn.disabled = true
     opacitySlider0.oninput()
@@ -167,18 +171,26 @@ async function main() {
     overlayVolume.hdr.scl_inter = 0
     overlayVolume.hdr.scl_slope = 1
     overlayVolume.img = new Uint8Array(img)
-    if (modelEntry.colormapPath) {
-      let cmap = await fetchJSON(modelEntry.colormapPath)
-      overlayVolume.setColormapLabel(cmap)
-      // n.b. most models create indexed labels, but those without colormap mask scalar input
-      overlayVolume.hdr.intent_code = 1002 // NIFTI_INTENT_LABEL
+    const isScalar = modelEntry.isScalar === true
+    if (isScalar) {
+      overlayVolume.hdr.scl_slope = 1 / 255
+      overlayVolume.colormap = "viridis"
     } else {
-      let colormap = opts.atlasSelectedColorTable.toLowerCase()
-      const cmaps = nv1.colormaps()
-      if (!cmaps.includes(colormap)) {
-        colormap = "actc"
+      
+      if (modelEntry.colormapPath) {
+        let cmap = await fetchJSON(modelEntry.colormapPath)
+        overlayVolume.setColormapLabel(cmap)
+        // n.b. most models create indexed labels, but those without colormap mask scalar input
+        overlayVolume.hdr.intent_code = 1002 // NIFTI_INTENT_LABEL
+      } else {
+        let colormap = opts.atlasSelectedColorTable.toLowerCase()
+        const cmaps = nv1.colormaps()
+        if (!cmaps.includes(colormap)) {
+          colormap = "actc"
+        }
+        overlayVolume.colormap = colormap
       }
-      overlayVolume.colormap = colormap
+
     }
     overlayVolume.opacity = opacitySlider1.value / 255
     await nv1.addVolume(overlayVolume)
@@ -243,12 +255,15 @@ async function main() {
     console.log(`Execution time: ${Math.round(performance.now() - startTime)} ms`)
   }
   async function applyFaster() {
-    const niiBuffer = await nv1.saveImage({volumeByIndex: nv1.volumes.length - 1}).buffer
+    const niiBuffer = await nv1.saveImage({volumeByIndex: nv1.volumes.length - 1})
     const niiFile = new File([niiBuffer], 'image.nii')
     let processor = niimath.image(niiFile)
     loadingCircle.classList.remove('hidden')
     //mesh with specified isosurface
-    const isoValue = 0.5
+    let isoValue = 0.5
+    if (nv1.volumes[nv1.volumes.length - 1].hdr.intent_code === 0) {
+      isoValue = 222 //isScalar
+    }
     //const largestCheckValue = largestCheck.checked
     let reduce = Math.min(Math.max(Number(shrinkPct.value) / 100, 0.01), 1)
     let hollowSz = Number(hollowSelect.value )