brainchop-webworker.js

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
} from './tensor-utils.js'

function callbackUI(message = '', progressFrac = -1, modalMessage = '', statData = []) {
  let statStr = []
  if (Object.keys(statData).length > 0) {
    function arrayToStr() {
      const list = {}
      for (const key in statData) {
        list[key] = statData[key]
      }
      return JSON.stringify(list)
    }
    statStr = arrayToStr(statData)
  }
  self.postMessage({
    cmd: 'ui',
    message,
    progressFrac,
    modalMessage,
    statData: statStr
  })
}

function callbackImg(img, opts, modelEntry) {
  self.postMessage({ cmd: 'img', img, opts, modelEntry })
}

async function inferenceFullVolumeSeqCovLayerPhase2(
  opts,
  modelEntry,
  model,
  slices_3d,
  num_of_slices,
  slice_height,
  slice_width,
  pipeline1_out,
  statData,
  niftiImage
) {
  // --Phase-2, After remove the skull try to allocate brain volume and make inferece

  console.log(' ---- Start FullVolume Inference with Sequential Conv Layer for phase-II ---- ')
  const quantileNorm = modelEntry.enableQuantileNorm
  if (quantileNorm) {
    // Quantile normalize function needs specific models to be used
    console.log('preModel Quantile normalization enabled')
    slices_3d = await quantileNormalizeVolumeData(slices_3d)
  } else {
    // Min Max Nomalize MRI data to be from 0 to 1
    console.log('preModel Min Max normalization enabled')
    slices_3d = await minMaxNormalizeVolumeData(slices_3d)
  }

  let mask_3d

  if (pipeline1_out == null) {
    // preModel is null

    // Check if thresholding the MRI to remove noisy voxels for better cropping is needed.
    const autoThresholdValue = modelEntry.autoThreshold

    if (autoThresholdValue > 0 && autoThresholdValue <= 1) {
      // Filtered MRI from noisy voxel below  autoThresholdValue
      mask_3d = await applyMriThreshold(slices_3d, autoThresholdValue)
    } else {
      console.log('No valid crop threshold value')
      // binarize original image
      mask_3d = await slices_3d.greater([0]).asType('bool')
    }
  } else {
    mask_3d = await pipeline1_out.greater([0]).asType('bool')
    // -- pipeline1_out.dispose()
  }

  console.log(' mask_3d shape :  ', mask_3d.shape)
  const [row_min, row_max, col_min, col_max, depth_min, depth_max] = await firstLastNonZero3D(mask_3d)
  mask_3d.dispose()
  // -- Reference voxel that cropped volume started slice with it
  const refVoxel = [row_min, col_min, depth_min]
  // -- Starting form refVoxel, size of bounding volume
  const boundVolSizeArr = [row_max - row_min + 1, col_max - col_min + 1, depth_max - depth_min + 1]

  // -- Extract 3d object (e.g. brain)
  const cropped_slices_3d = await slices_3d.slice(
    [row_min, col_min, depth_min],
    [row_max - row_min + 1, col_max - col_min + 1, depth_max - depth_min + 1]
  )
  slices_3d.dispose()

  // -- Padding size add to cropped brain
  const pad = modelEntry.cropPadding

  // Create margin around the bounding volume
  let cropped_slices_3d_w_pad = await addZeroPaddingTo3dTensor(cropped_slices_3d, [pad, pad], [pad, pad], [pad, pad])
  console.log(' cropped slices_3d with padding shape:  ', cropped_slices_3d_w_pad.shape)

  cropped_slices_3d.dispose()

  if (opts.drawBoundingVolume) {
    let testVol = await removeZeroPaddingFrom3dTensor(cropped_slices_3d_w_pad, pad, pad, pad)
    console.log(' outLabelVolume without padding shape :  ', testVol.shape)

    testVol = await resizeWithZeroPadding(testVol, num_of_slices, slice_height, slice_width, refVoxel, boundVolSizeArr)
    console.log(' outLabelVolume final shape after resizing :  ', testVol.shape)
    draw3dObjBoundingVolume(tf.unstack(testVol), opts, modelEntry, callbackImg)
    testVol.dispose()

    return 0
  }

  statData.Brainchop_Ver = 'FullVolume'
  const res = await model
  try {
    let startTime = performance.now()
    const inferenceStartTime = performance.now()
    // maxLabelPredicted in whole volume of the brain
    let maxLabelPredicted = 0
    const transpose = modelEntry.enableTranspose

    if (transpose) {
      cropped_slices_3d_w_pad = await cropped_slices_3d_w_pad.transpose()
      console.log('Input transposed for pre-model')
    } else {
      console.log('Transpose not enabled for pre-model')
    }

    let i = 1
    const layersLength = res.layers.length
    console.log('res.layers.length ', layersLength)

    const isChannelLast = isModelChnlLast(res)
    const batchSize = opts.batchSize
    const numOfChan = opts.numOfChan
    let adjusted_input_shape
    // -- Adjust model input shape
    if (isChannelLast) {
      res.layers[0].batchInputShape[1] = cropped_slices_3d_w_pad.shape[0]
      res.layers[0].batchInputShape[2] = cropped_slices_3d_w_pad.shape[1]
      res.layers[0].batchInputShape[3] = cropped_slices_3d_w_pad.shape[2]

      adjusted_input_shape = [
        batchSize,
        res.layers[0].batchInputShape[1],
        res.layers[0].batchInputShape[2],
        res.layers[0].batchInputShape[3],
        numOfChan
      ]
    } else {
      res.layers[0].batchInputShape[2] = cropped_slices_3d_w_pad.shape[0]
      res.layers[0].batchInputShape[3] = cropped_slices_3d_w_pad.shape[1]
      res.layers[0].batchInputShape[4] = cropped_slices_3d_w_pad.shape[2]

      adjusted_input_shape = [
        batchSize,
        numOfChan,
        res.layers[0].batchInputShape[2],
        res.layers[0].batchInputShape[3],
        res.layers[0].batchInputShape[4]
      ]
    }

    console.log(' Model batch input shape : ', res.layers[0].batchInputShape)
    // -- batchInputShape {Array} input_shape - e.g. [?, D, H, W, Ch] or [?, Ch, D, H, W]

    statData.Input_Shape = JSON.stringify(res.layers[0].batchInputShape)
    statData.Output_Shape = JSON.stringify(res.output.shape)
    statData.Channel_Last = await isChannelLast
    statData.Model_Param = await getModelNumParameters(res)
    statData.Model_Layers = await getModelNumLayers(res)
    statData.Model = modelEntry.modelName
    statData.Seq_Conv = modelEntry.enableSeqConv
    // statData.Extra_Info = null

    // Determine the number of output channels in the last layer of the model
    //  e.g. 3, 50, 104
    const outputLayer = res.layers[res.layers.length - 1]
    console.log('Output Layer : ', outputLayer)

    const expected_Num_labels = isChannelLast
      ? outputLayer.outputShape[outputLayer.outputShape.length - 1]
      : outputLayer.outputShape[1]
    console.log('Num of output channels x: ', expected_Num_labels)

    const curTensor = []
    curTensor[0] = await cropped_slices_3d_w_pad.reshape(adjusted_input_shape)
    while (true) {
      try {
        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
        }

        tf.dispose(curTensor[i - 1])
      } catch (err) {
        const errTxt = 'Your graphics card (e.g. Intel) may not be compatible with WebGL. ' + err.message
        callbackUI(errTxt, -1, errTxt)

        tf.engine().endScope()
        tf.engine().disposeVariables()

        statData.Inference_t = Infinity
        statData.Postprocess_t = Infinity
        statData.Status = 'Fail'
        statData.Error_Type = err.message
        statData.Extra_Err_Info = 'Failed while model layer ' + i + ' apply'

        callbackUI('', -1, '', statData)

        return 0
      }

      console.log('layer output Tensor shape : ', curTensor[i].shape)
      console.log('layer count params ', res.layers[i].countParams())

      res.layers[i].dispose()
      curTensor[i - 1].dispose()

      callbackUI('Layer ' + i.toString(), (i + 1) / layersLength)
      if (tf.memory().unreliable) {
        const unreliableReasons = 'unreliable reasons :' + tf.memory().reasons
        callbackUI(unreliableReasons, NaN, unreliableReasons)
      }
      if (i === layersLength - 2) {
        // Stop before the last layer or classification layer.

        // // Create an instance of SequentialConvLayer
        // The second parameter is important for memory,
        // the larger it is, the more memory it uses
        // it was 8, but I set it to 3, got a different error
        // let seqConvLayer = new SequentialConvLayer(res, 10, isChannelLast)
        const seqConvLayer = await new SequentialConvLayer(res, 10, isChannelLast, callbackUI)

        // Apply the last output tensor to the seq. instance
        let outputTensor = null
        const profileInfo = await tf.profile(async () => {
          // Your tensor operations here
          outputTensor = await seqConvLayer.apply(curTensor[i])
        })
        console.log('profileInfo : ', profileInfo)

        // -- document.getElementById("progressBarChild").style.width = 0 + "%";

        // Dispose the previous layer input tensor
        tf.dispose(curTensor[i])
        // delete the used class
        // ? delete seqConvLayer

        // You can now use 'outputTensor' as needed
        console.log(' Output tensor', outputTensor)
        console.log(' Output tensor shape : ', outputTensor.shape)
        // Array(3) [ 256, 256, 256 ]

        if (outputTensor.shape.length !== 3) {
          const msg = 'Output tensor shape should be 3 dims but it is ' + outputTensor.shape.length
          callbackUI(msg, -1, msg)
        }

        const Inference_t = ((performance.now() - startTime) / 1000).toFixed(4)

        console.log(' find array max ')
        const curBatchMaxLabel = await outputTensor.max().dataSync()[0]
        if (maxLabelPredicted < curBatchMaxLabel) {
          maxLabelPredicted = curBatchMaxLabel
        }

        const numSegClasses = maxLabelPredicted + 1
        console.log('Predicted num of segmentation classes', numSegClasses)
        statData.Actual_Labels = numSegClasses
        statData.Expect_Labels = expected_Num_labels
        statData.NumLabels_Match = numSegClasses === expected_Num_labels
        if (numSegClasses !== expected_Num_labels) {
          const msg = 'expected ' + expected_Num_labels + ' labels, but the predicted are ' + numSegClasses
          callbackUI(msg, -1, msg)
        }

        // -- Transpose back to original unpadded size
        let outLabelVolume = outputTensor.reshape([
          cropped_slices_3d_w_pad.shape[0],
          cropped_slices_3d_w_pad.shape[1],
          cropped_slices_3d_w_pad.shape[2]
        ])
        tf.dispose(outputTensor)

        // Transpose MRI data to be match pytorch/keras input output
        if (transpose) {
          console.log('outLabelVolume transposed')
          outLabelVolume = outLabelVolume.transpose()
        }

        outLabelVolume = await removeZeroPaddingFrom3dTensor(outLabelVolume, pad, pad, pad)
        console.log(' outLabelVolume without padding shape :  ', outLabelVolume.shape)
        outLabelVolume = await resizeWithZeroPadding(
          outLabelVolume,
          num_of_slices,
          slice_height,
          slice_width,
          refVoxel,
          boundVolSizeArr
        )
        console.log(' outLabelVolume final shape after resizing :  ', outLabelVolume.shape)

        // let filterOutWithPreMask =  inferenceModelsList[$$("selectModel").getValue() - 1]["filterOutWithPreMask"]
        const filterOutWithPreMask = modelEntry.filterOutWithPreMask
        // To clean the skull area wrongly segmented inphase-2.
        if (pipeline1_out != null && opts.isBrainCropMaskBased && filterOutWithPreMask) {
          const bin = await binarizeVolumeDataTensor(pipeline1_out)
          outLabelVolume = await outLabelVolume.mul(bin)
        }

        startTime = performance.now()
        // Generate output volume or slices
        console.log('Generating correct output')
        let outimg
        try {
          const img = await new Uint32Array(outLabelVolume.dataSync())
          const Vshape = outLabelVolume.shape
          const Vtype = outLabelVolume.dtype
          outimg = await generateOutputSlicesV2(
            img,
            Vshape,
            Vtype,
            num_of_slices,
            numSegClasses,
            slice_height,
            slice_width,
            modelEntry,
            opts,
            niftiImage
          )
          console.log(' Phase-2 num of tensors after generateOutputSlicesV2: ', tf.memory().numTensors)

          tf.dispose(outLabelVolume)
          tf.engine().endScope()
          tf.engine().disposeVariables()
        } catch (error) {
          // -- Timing data to collect
          tf.engine().endScope()
          tf.engine().disposeVariables()
          console.log('Error while generating output: ', error)
          const msg = 'Failed while generating output due to limited browser memory available'
          callbackUI(msg, -1, msg)

          statData.Inference_t = Inference_t
          statData.Postprocess_t = Infinity
          statData.Status = 'Fail'
          statData.Error_Type = error.message
          statData.Extra_Err_Info = 'Failed while generating output'

          callbackUI('', -1, '', statData)

          return 0
        }
        const Postprocess_t = ((performance.now() - startTime) / 1000).toFixed(4)

        console.log(
          'Processing the whole brain volume in tfjs for multi-class output mask took : ',
          ((performance.now() - inferenceStartTime) / 1000).toFixed(4) + '  Seconds'
        )

        // -- Timing data to collect
        statData.Inference_t = Inference_t
        statData.Postprocess_t = Postprocess_t
        statData.Status = 'OK'

        callbackUI('', -1, '', statData)
        callbackUI('Segmentation finished', 0)
        callbackImg(outimg, opts, modelEntry)
        return 0
      } else {
        i++
      }
    }
  } catch (err) {
    callbackUI(err.message, -1, err.message)
    console.log(
      'If webgl context is lost, try to restore webgl context by visit the link ' +
        '<a href="https://support.biodigital.com/hc/en-us/articles/218322977-How-to-turn-on-WebGL-in-my-browser">here</a>'
    )
    if (tf.memory().unreliable) {
      const unreliableReasons = 'unreliable reasons :' + tf.memory().reasons
      callbackUI(unreliableReasons, NaN, unreliableReasons)
    }
  }
}

async function inferenceFullVolumePhase2(
  model,
  slices_3d,
  num_of_slices,
  slice_height,
  slice_width,
  pipeline1_out,
  modelEntry,
  statData,
  opts,
  niftiImage
) {
  let outimg = []
  // --Phase-2, After remove the skull try to allocate brain volume and make inferece
  console.log(' ---- Start FullVolume inference phase-II ---- ')
  const quantileNorm = modelEntry.enableQuantileNorm
  if (quantileNorm) {
    // Quantile normalize function needs specific models to be used
    console.log('preModel Quantile normalization enabled')
    slices_3d = await quantileNormalizeVolumeData(slices_3d)
  } else {
    // Min Max Nomalize MRI data to be from 0 to 1
    console.log('preModel Min Max normalization enabled')
    slices_3d = await minMaxNormalizeVolumeData(slices_3d)
  }
  let mask_3d
  if (pipeline1_out == null) {
    // preModel is null

    // Check if thresholding the MRI to remove noisy voxels for better cropping is needed.
    const autoThresholdValue = modelEntry.autoThreshold

    if (autoThresholdValue > 0 && autoThresholdValue <= 1) {
      // Filtered MRI from noisy voxel below  autoThresholdValue
      mask_3d = await applyMriThreshold(slices_3d, autoThresholdValue)
    } else {
      console.log('No valid crop threshold value')
      // binarize original image
      mask_3d = await slices_3d.greater([0]).asType('bool')
    }
  } else {
    mask_3d = await pipeline1_out.greater([0]).asType('bool')
    // -- pipeline1_out.dispose()
  }
  console.log(' mask_3d shape :  ', mask_3d.shape)
  const [row_min, row_max, col_min, col_max, depth_min, depth_max] = await firstLastNonZero3D(mask_3d)
  mask_3d.dispose()
  // -- Reference voxel that cropped volume started slice with it
  const refVoxel = [row_min, col_min, depth_min]
  console.log('refVoxel :', refVoxel)

  // -- Starting form refVoxel, size of bounding volume
  const boundVolSizeArr = [row_max - row_min + 1, col_max - col_min + 1, depth_max - depth_min + 1]

  console.log('boundVolSizeArr :', boundVolSizeArr)
  // -- Extract 3d object (e.g. brain)
  const cropped_slices_3d = slices_3d.slice(
    [row_min, col_min, depth_min],
    [row_max - row_min + 1, col_max - col_min + 1, depth_max - depth_min + 1]
  )

  slices_3d.dispose()

  // -- Padding size add to cropped brain
  const pad = modelEntry.cropPadding

  // Create margin around the bounding volume
  let cropped_slices_3d_w_pad = await addZeroPaddingTo3dTensor(cropped_slices_3d, [pad, pad], [pad, pad], [pad, pad])
  console.log(' cropped slices_3d with padding shape:  ', cropped_slices_3d_w_pad.shape)

  cropped_slices_3d.dispose()

  // -- Test dim after padding ..
  // for (let i = 0; i < cropped_slices_3d_w_pad.rank; i++) {
  //     if(cropped_slices_3d_w_pad.shape[i] > 256) {
  //          console.log(" cropped_slices_3d_w_pad > 256 ")
  //     }

  // }

  if (opts.drawBoundingVolume) {
    let testVol = await removeZeroPaddingFrom3dTensor(cropped_slices_3d_w_pad, pad, pad, pad)
    console.log(' outLabelVolume without padding shape :  ', testVol.shape)

    testVol = await resizeWithZeroPadding(testVol, num_of_slices, slice_height, slice_width, refVoxel, boundVolSizeArr)
    console.log(' outLabelVolume final shape after resizing :  ', testVol.shape)
    draw3dObjBoundingVolume(tf.unstack(testVol), opts, modelEntry, callbackImg)
    testVol.dispose()

    return 0
  }

  statData.Brainchop_Ver = 'FullVolume'
  let startTime = performance.now()
  let adjusted_input_shape = []
  const res = await model
  try {
    startTime = performance.now()
    const inferenceStartTime = performance.now()
    // maxLabelPredicted in whole volume of the brain
    let maxLabelPredicted = 0
    const transpose = modelEntry.enableTranspose

    if (transpose) {
      cropped_slices_3d_w_pad = cropped_slices_3d_w_pad.transpose()
      console.log('Input transposed for pre-model')
    } else {
      console.log('Transpose not enabled for pre-model')
    }

    let i = 1
    const layersLength = res.layers.length
    console.log('res.layers.length ', layersLength)

    const isChannelLast = isModelChnlLast(res)
    const batchSize = opts.batchSize
    const numOfChan = opts.numOfChan

    // -- Adjust model input shape
    if (isChannelLast) {
      res.layers[0].batchInputShape[1] = cropped_slices_3d_w_pad.shape[0]
      res.layers[0].batchInputShape[2] = cropped_slices_3d_w_pad.shape[1]
      res.layers[0].batchInputShape[3] = cropped_slices_3d_w_pad.shape[2]

      adjusted_input_shape = [
        batchSize,
        res.layers[0].batchInputShape[1],
        res.layers[0].batchInputShape[2],
        res.layers[0].batchInputShape[3],
        numOfChan
      ]
    } else {
      res.layers[0].batchInputShape[2] = cropped_slices_3d_w_pad.shape[0]
      res.layers[0].batchInputShape[3] = cropped_slices_3d_w_pad.shape[1]
      res.layers[0].batchInputShape[4] = cropped_slices_3d_w_pad.shape[2]

      adjusted_input_shape = [
        batchSize,
        numOfChan,
        res.layers[0].batchInputShape[2],
        res.layers[0].batchInputShape[3],
        res.layers[0].batchInputShape[4]
      ]
    }

    console.log(' Model batch input shape : ', res.layers[0].batchInputShape)
    // -- batchInputShape {Array} input_shape - e.g. [?, D, H, W, Ch] or [?, Ch, D, H, W]

    statData.Input_Shape = JSON.stringify(res.layers[0].batchInputShape)
    statData.Output_Shape = JSON.stringify(res.output.shape)
    statData.Channel_Last = await isChannelLast
    statData.Model_Param = await getModelNumParameters(res)
    statData.Model_Layers = await getModelNumLayers(res)
    statData.Model = modelEntry.modelName
    // statData.Extra_Info = null

    const curTensor = []
    curTensor[0] = cropped_slices_3d_w_pad.reshape(adjusted_input_shape)
    // console.log("curTensor[0] :", curTensor[0].dataSync())

    while (true) {
      try {
        // -- curTensor[i] = res.layers[i].apply( curTensor[i-1])
        curTensor[i] = res.layers[i].apply(curTensor[i - 1])
      } catch (err) {
        callbackUI(err.message, -1, err.message)
        tf.engine().endScope()
        tf.engine().disposeVariables()

        statData.Inference_t = Infinity
        statData.Postprocess_t = Infinity
        statData.Status = 'Fail'
        statData.Error_Type = err.message
        statData.Extra_Err_Info = 'Failed while model layer ' + i + ' apply'

        callbackUI('', -1, '', statData)

        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())
      res.layers[i].dispose()
      curTensor[i - 1].dispose()
      if (tf.memory().unreliable) {
        const unreliableReasons = 'unreliable reasons :' + tf.memory().reasons
        callbackUI(unreliableReasons, NaN, unreliableReasons)
      }

      if (i === layersLength - 1) {
        // prediction = res.layers[res.layers.length-1].apply(curTensor[i])
        // curTensor[i].print()
        // outputDataBeforArgmx = Array.from(curTensor[i].dataSync())

        const axis = isChannelLast ? -1 : 1
        console.log(' find argmax ')
        console.log('last Tensor shape : ', curTensor[i].shape)
        // -- curTensor[i].shape  e.g. [ 1, 256, 256, 256, 3 ]
        const expected_Num_labels = isChannelLast ? curTensor[i].shape[4] : curTensor[i].shape[1]
        let prediction_argmax

        // Try for argMax with model output tensor.

        try {
          const argMaxTime = performance.now()
          console.log(' Try tf.argMax for fullVolume ..')
          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
          if (axis === -1) {
            try {
              const argMaxLargeTime = performance.now()
              console.log(' tf.argMax failed .. try argMaxLarge ..')
              callbackUI('', -1, 'tensor2LightBuffer() is not dead code?')
              callbackUI('', -1, 'argMaxLarge() is not dead code?')
              console.log(
                'argMaxLarge for fullVolume takes : ',
                ((performance.now() - argMaxLargeTime) / 1000).toFixed(4)
              )
            } catch (err2) {
              const errTxt = "argMax buffer couldn't be created due to limited memory resources."
              callbackUI(errTxt, -1, errTxt)

              tf.engine().endScope()
              tf.engine().disposeVariables()

              statData.Inference_t = Infinity
              statData.Postprocess_t = Infinity
              statData.Status = 'Fail'
              statData.Error_Type = err2.message
              statData.Extra_Err_Info = 'prediction_argmax from argMaxLarge failed'

              callbackUI('', -1, '', statData)
              return 0
            }
          } else {
            // if channel first ..
            const errTxt = "argMax buffer couldn't be created due to limited memory resources."
            callbackUI(errTxt, -1, errTxt)

            prediction_argmax.dispose()

            tf.engine().endScope()
            tf.engine().disposeVariables()

            statData.Inference_t = Infinity
            statData.Postprocess_t = Infinity
            statData.Status = 'Fail'
            statData.Error_Type = err1.message
            statData.Extra_Err_Info = 'prediction_argmax from argMaxLarge not support yet channel first'

            callbackUI('', -1, '', statData)

            return 0
          }
        }

        console.log(' prediction_argmax shape : ', prediction_argmax.shape)
        // -- prediction_argmax.shape  : [ 1, 256, 256, 256]

        const Inference_t = ((performance.now() - startTime) / 1000).toFixed(4)

        // outputDataBeforArgmx = Array.from(prediction_argmax.dataSync())
        tf.dispose(curTensor[i])
        console.log(' find array max ')
        const curBatchMaxLabel = await prediction_argmax.max().dataSync()[0]

        if (maxLabelPredicted < curBatchMaxLabel) {
          maxLabelPredicted = curBatchMaxLabel
        }

        const numSegClasses = maxLabelPredicted + 1
        console.log('numSegClasses', numSegClasses)
        statData.Actual_Labels = numSegClasses
        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)
        }

        // -- Transpose back to original unpadded size
        let outLabelVolume = prediction_argmax.reshape([
          cropped_slices_3d_w_pad.shape[0],
          cropped_slices_3d_w_pad.shape[1],
          cropped_slices_3d_w_pad.shape[2]
        ])
        tf.dispose(prediction_argmax)

        // Transpose MRI data to be match pytorch/keras input output
        if (transpose) {
          console.log('outLabelVolume transposed')
          outLabelVolume = outLabelVolume.transpose()
        }
        outLabelVolume = await removeZeroPaddingFrom3dTensor(outLabelVolume, pad, pad, pad)
        console.log(' outLabelVolume without padding shape :  ', outLabelVolume.shape)
        outLabelVolume = await resizeWithZeroPadding(
          outLabelVolume,
          num_of_slices,
          slice_height,
          slice_width,
          refVoxel,
          boundVolSizeArr
        )
        console.log(' outLabelVolume final shape after resizing :  ', outLabelVolume.shape)

        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)
        }

        startTime = performance.now()
        // Generate output volume or slices
        console.log('Generating correct output')

        try {
          const img = new Uint32Array(outLabelVolume.dataSync())
          const Vshape = outLabelVolume.shape
          const Vtype = outLabelVolume.dtype
          tf.dispose(outLabelVolume)
          tf.engine().endScope()
          tf.engine().disposeVariables()
          outimg = await generateOutputSlicesV2(
            img,
            Vshape,
            Vtype,
            num_of_slices,
            numSegClasses,
            slice_height,
            slice_width,
            modelEntry,
            opts,
            niftiImage
          )
          console.log(' Phase-2 num of tensors after generateOutputSlicesV2: ', tf.memory().numTensors)
        } catch (error) {
          // -- Timing data to collect
          tf.engine().endScope()
          tf.engine().disposeVariables()

          const errTxt = 'Failed while generating output due to limited browser memory available'
          callbackUI(errTxt, -1, errTxt)
          statData.Inference_t = Inference_t
          statData.Postprocess_t = Infinity
          statData.Status = 'Fail'
          statData.Error_Type = error.message
          statData.Extra_Err_Info = 'Failed while generating output'

          callbackUI('', -1, '', statData)

          return 0
        }

        const Postprocess_t = ((performance.now() - startTime) / 1000).toFixed(4)

        tf.engine().disposeVariables()

        console.log(
          'Processing the whole brain volume in tfjs for multi-class output mask took : ',
          ((performance.now() - inferenceStartTime) / 1000).toFixed(4) + '  Seconds'
        )

        // -- Timing data to collect
        statData.Inference_t = Inference_t
        statData.Postprocess_t = Postprocess_t
        statData.Status = 'OK'
        callbackUI('Segmentation finished', 0)
        callbackUI('', -1, '', statData)
        callbackImg(outimg, opts, modelEntry)

        return 0
      }
      i++
    }
  } catch (err) {
    callbackUI(err.message, -1, err.message)
    console.log(
      'If webgl context is lost, try to restore webgl context by visit the link ' +
        '<a href="https://support.biodigital.com/hc/en-us/articles/218322977-How-to-turn-on-WebGL-in-my-browser">here</a>'
    )
  }
}

async function inferenceFullVolumePhase1(
  model,
  slices_3d,
  num_of_slices,
  slice_height,
  slice_width,
  isModelFullVol,
  modelEntry,
  statData,
  opts,
  niftiHeader,
  niftiImage
) {
  statData.No_SubVolumes = 1
  // load pre-model for inference first, can be null if no pre-model such as GWM models
  if (modelEntry.preModelId) {
    const preModel = await load_model(opts.rootURL + inferenceModelsList[modelEntry.preModelId - 1].path)
    const transpose = inferenceModelsList[modelEntry.preModelId - 1].enableTranspose
    const quantileNorm = inferenceModelsList[modelEntry.preModelId - 1].enableQuantileNorm
    let preModel_slices_3d = null

    // -- If pre-model is not null then slices_3d mask will be generated..
    // -- The mask is needed to remove the skull and set noise in background to 0, and get the brain bounding volume properly
    const slices_3d_mask = null

    if (quantileNorm) {
      // Quantile normalize function needs specific models to be used
      console.log('preModel Quantile normalization enabled')
      preModel_slices_3d = await quantileNormalizeVolumeData(slices_3d)
    } else {
      // Min Max Nomalize MRI data to be from 0 to 1
      console.log('preModel Min Max normalization enabled')
      preModel_slices_3d = await minMaxNormalizeVolumeData(slices_3d)
    }

    // -- Transpose MRI data to be match pytorch/keras input output
    // -- Check if pre-model needs transpose..
    if (transpose) {
      preModel_slices_3d = preModel_slices_3d.transpose()
      console.log('Input transposed for pre-model')
    } else {
      console.log('Transpose not enabled for pre-model')
    }

    statData.Brainchop_Ver = 'PreModel_FV' // e.g. "PreModel_FV"

    // preModel.then(function (res) {
    const res = await preModel

    try {
      const inferenceStartTime = performance.now()
      const preModelObject = res

      // read input shape from model.json object
      const preModelBatchInputShape = preModelObject.layers[0].batchInputShape
      console.log(' Pre-Model batch input shape : ', preModelBatchInputShape)

      // -- Verify input shape
      if (preModelBatchInputShape.length !== 5) {
        const errTxt = 'The pre-model input shape must be 5D '
        callbackUI(errTxt, -1, errTxt)
        return 0
      }

      const isPreModelChannelLast = await isModelChnlLast(preModelObject)
      const batchSize = opts.batchSize
      const numOfChan = opts.numOfChan
      let batch_D, batch_H, batch_W
      let preModel_input_shape
      if (isPreModelChannelLast) {
        console.log('Pre-Model Channel Last')
        if (isNaN(preModelBatchInputShape[4]) || preModelBatchInputShape[4] !== 1) {
          const errTxt = 'The number of channels for pre-model input shape must be 1'
          callbackUI(errTxt, -1, errTxt)
          return 0
        }

        batch_D = preModelBatchInputShape[1]
        batch_H = preModelBatchInputShape[2]
        batch_W = preModelBatchInputShape[3]

        preModel_input_shape = [batchSize, batch_D, batch_H, batch_W, numOfChan]
      } else {
        console.log('Pre-Model Channel First')
        if (isNaN(preModelBatchInputShape[1]) || preModelBatchInputShape[1] !== 1) {
          const errTxt = 'The number of channels for pre-model input shape must be 1'
          callbackUI(errTxt, -1, errTxt)
          return 0
        }

        batch_D = preModelBatchInputShape[2]
        batch_H = preModelBatchInputShape[3]
        batch_W = preModelBatchInputShape[4]

        preModel_input_shape = [batchSize, numOfChan, batch_D, batch_H, batch_W]
      }

      statData.Input_Shape = JSON.stringify(preModel_input_shape)
      statData.Output_Shape = JSON.stringify(preModelObject.output.shape)
      statData.Channel_Last = await isPreModelChannelLast
      statData.Model_Param = await getModelNumParameters(preModelObject)
      statData.Model_Layers = await getModelNumLayers(preModelObject)

      // maxLabelPredicted in whole volume of the brain
      let maxLabelPredicted = 0

      let i = 1
      const layersLength = res.layers.length

      const curTensor = []
      // -- reshape MRI to model input shape
      curTensor[0] = preModel_slices_3d.reshape(preModel_input_shape)

      // Dispose the volume
      tf.dispose(preModel_slices_3d)
      while (true) {
        try {
          curTensor[i] = res.layers[i].apply(curTensor[i - 1])
        } catch (err) {
          const errTxt = 'Your graphics card (e.g. Intel) may not be compatible with WebGL. ' + err.message
          callbackUI(errTxt, -1, errTxt)

          tf.engine().endScope()
          tf.engine().disposeVariables()

          statData.Inference_t = Infinity
          statData.Postprocess_t = Infinity
          statData.Status = 'Fail'
          statData.Error_Type = err.message
          statData.Extra_Err_Info = 'PreModel Failed while model layer ' + i + ' apply'

          callbackUI('', -1, '', statData)

          return 0
        }

        res.layers[i].dispose()
        curTensor[i - 1].dispose()

        callbackUI('Layer ' + i.toString(), (i + 1) / layersLength)
        if (tf.memory().unreliable) {
          const unreliableReasons = 'unreliable reasons :' + tf.memory().reasons
          callbackUI(unreliableReasons, NaN, unreliableReasons)
        }

        if (i === layersLength - 1) {
          // -- prediction = res.layers[res.layers.length-1].apply(curTensor[i])
          // -- curTensor[i].print()
          // -- outputDataBeforArgmx = Array.from(curTensor[i].dataSync())

          const axis = isPreModelChannelLast ? -1 : 1
          console.log(' find argmax ')
          console.log('last Tensor shape : ', curTensor[i].shape)
          // -- curTensor[i].shape  : [ 1, 256, 256, 256, 3 ]
          const expected_Num_labels = isPreModelChannelLast ? curTensor[i].shape[4] : curTensor[i].shape[1]
          let prediction_argmax

          // Try for argMax with model output tensor.

          try {
            console.log(' Try tf.argMax for fullVolume ..')
            prediction_argmax = await tf.argMax(curTensor[i], axis)
          } catch (err1) {
            // if channel last
            if (axis === -1) {
              try {
                const argMaxLargeTime = performance.now()
                console.log(' tf.argMax failed .. try argMaxLarge ..')
                callbackUI('', -1, 'tensor2LightBuffer() is not dead code?')
                callbackUI('', -1, 'argMaxLarge() is not dead code?')
                console.log(
                  'argMaxLarge for fullVolume takes : ',
                  ((performance.now() - argMaxLargeTime) / 1000).toFixed(4)
                )
              } catch (err2) {
                const errTxt = "argMax buffer couldn't be created due to limited memory resources."
                callbackUI(errTxt, -1, errTxt)

                prediction_argmax.dispose()

                tf.engine().endScope()
                tf.engine().disposeVariables()

                statData.Inference_t = Infinity
                statData.Postprocess_t = Infinity
                statData.Status = 'Fail'
                statData.Error_Type = err2.message
                statData.Extra_Err_Info = 'preModel prediction_argmax from argMaxLarge failed'

                callbackUI('', -1, '', statData)

                return 0
              }
            } else {
              // if channel first ..
              const errTxt = "argMax buffer couldn't be created due to limited memory resources."
              callbackUI(errTxt, -1, errTxt)

              prediction_argmax.dispose()

              tf.engine().endScope()
              tf.engine().disposeVariables()

              statData.Inference_t = Infinity
              statData.Postprocess_t = Infinity
              statData.Status = 'Fail'
              statData.Error_Type = err1.message
              statData.Extra_Err_Info = 'preModel prediction_argmax from argMaxLarge not support yet channel first'

              callbackUI('', -1, '', statData)

              return 0
            }
          }

          console.log(' Pre-model prediction_argmax shape : ', prediction_argmax.shape)
          // -- prediction_argmax.shape  : [ 1, 256, 256, 256]

          const Inference_t = ((performance.now() - inferenceStartTime) / 1000).toFixed(4)

          tf.dispose(curTensor[i])

          console.log(' Pre-model find array max ')
          const curBatchMaxLabel = await prediction_argmax.max().dataSync()[0]

          if (maxLabelPredicted < curBatchMaxLabel) {
            maxLabelPredicted = curBatchMaxLabel
          }

          const numSegClasses = maxLabelPredicted + 1
          console.log('Pre-model numSegClasses', numSegClasses)

          statData.Actual_Labels = numSegClasses
          statData.Expect_Labels = expected_Num_labels
          statData.NumLabels_Match = numSegClasses === expected_Num_labels

          // -- Transpose back to original unpadded size
          let outLabelVolume = await prediction_argmax.reshape([num_of_slices, slice_height, slice_width])
          tf.dispose(prediction_argmax)
          // Transpose MRI data to be match pytorch/keras input output
          if (transpose) {
            console.log('Pre-model outLabelVolume transposed')
            outLabelVolume = outLabelVolume.transpose()
          }
          const startTime = performance.now()
          // Generate output volume or slices
          console.log('Generating pre-model output')
          let slices_3d_mask
          try {
            const unstackOutVolumeTensor = await tf.unstack(outLabelVolume)
            slices_3d_mask = await generateBrainMask(
              unstackOutVolumeTensor,
              num_of_slices,
              slice_height,
              slice_width,
              modelEntry,
              opts,
              niftiHeader,
              niftiImage,
              false
            )
            await tf.dispose(outLabelVolume)
            console.log(' Phase-1 num of tensors after generateBrainMask: ', tf.memory().numTensors)
          } catch (error) {
            // -- Timing data to collect
            tf.engine().endScope()
            tf.engine().disposeVariables()

            const errTxt = 'Failed while generating pre-model output due to limited browser memory available'
            callbackUI(errTxt, -1, errTxt)

            statData.Inference_t = Inference_t
            statData.Postprocess_t = Infinity
            statData.Status = 'Fail'
            statData.Error_Type = error.message
            statData.Extra_Err_Info = 'Pre-model failed while generating output'

            callbackUI('', -1, '', statData)

            return 0
          }
          const Postprocess_t = ((performance.now() - startTime) / 1000).toFixed(4)
          console.log(
            'Pre-model processing the whole brain volume in tfjs tooks for multi-class output mask : ',
            ((performance.now() - inferenceStartTime) / 1000).toFixed(4) + '  Seconds'
          )

          // -- Timing data to collect
          statData.Inference_t = Inference_t
          statData.Postprocess_t = Postprocess_t
          statData.Status = 'OK'

          callbackUI('', -1, '', statData)

          if (slices_3d_mask == null) {
            const msg = 'slice_3d_mask failed ...'
            callbackUI(msg, -1, msg)
            return 0
          } else {
            // --Phase-2, After remove the skull try to allocate brain volume and make inferece
            console.log('--- pre-model done ---')
            // --mask_3d = slices_3d_mask.greater([0]).asType('bool')
            // --slices_3d_mask.dispose()

            if (isModelFullVol) {
              if (modelEntry.enableSeqConv) {
                // Mask cropping & seq conv
                // Non-Atlas model (e.g. GWM) needs sequential convolution layer.
                // Sequential convolution layer to be used after cropping - slow but reliable on most machines
                console.log('------ Mask Cropping & Seq Convoluton ------')
                await inferenceFullVolumeSeqCovLayerPhase2(
                  opts,
                  modelEntry,
                  model,
                  slices_3d,
                  num_of_slices,
                  slice_height,
                  slice_width,
                  slices_3d_mask,
                  statData,
                  niftiImage
                )
                return 0
                // inferenceFullVolumeSeqCovLayerPhase2(model, slices_3d.transpose(), num_of_slices, slice_height, slice_width, slices_3d_mask)
              } else {
                // Mask cropping BUT no seq conv
                console.log('------ Mask Cropping  -  NO Seq Convoluton ------')
                await inferenceFullVolumePhase2(
                  model,
                  slices_3d,
                  num_of_slices,
                  slice_height,
                  slice_width,
                  slices_3d_mask,
                  modelEntry,
                  statData,
                  opts,
                  niftiImage
                )
                // inferenceFullVolumePhase2(model, slices_3d.transpose(), num_of_slices, slice_height, slice_width, slices_3d_mask)
              }
            } else {
              // -- In version 3.0.0 this function not used
              callbackUI('', -1, 'inferenceSubVolumes() is not dead code?')
            }
          }
        }
        i++
      }
    } catch (err) {
      callbackUI(err.message, -1, err.message)
      console.log(
        'If webgl context is lost, try to restore webgl context by visit the link ' +
          '<a href="https://support.biodigital.com/hc/en-us/articles/218322977-How-to-turn-on-WebGL-in-my-browser">here</a>'
      )

      // document.getElementById("webGl2Status").style.backgroundColor =  isWebGL2ContextLost() ? "Red" : "Green"
      // document.getElementById("memoryStatus").style.backgroundColor =  tf.memory().unreliable ? "Red" : "Green"
    }
    // })

    // -- if(...)  end
  } else {
    // No preModel

    // --Phase-2, After remove the skull try to allocate brain volume and make inferece
    console.log('--- No pre-model is selected ---')
    console.log('------ Run voxel cropping ------')
    // -- mask_3d = slices_3d.greater([0]).asType('bool')

    if (isModelFullVol) {
      if (modelEntry.enableSeqConv) {
        // Voxel cropping & seq conv
        // Non-Atlas model (e.g. GWM) needs sequential convolution layer.
        // Sequential convolution layer to be used after cropping - slow but reliable on most machines
        console.log('------ Seq Convoluton ------')
        await inferenceFullVolumeSeqCovLayerPhase2(
          opts,
          modelEntry,
          model,
          slices_3d,
          num_of_slices,
          slice_height,
          slice_width,
          null,
          statData,
          niftiImage
        )
      } else {
        // Voxel cropping BUT no seq conv
        // todo: we do not use result const outimg = await
        inferenceFullVolumePhase2(
          model,
          slices_3d,
          num_of_slices,
          slice_height,
          slice_width,
          null,
          modelEntry,
          statData,
          opts,
          niftiImage
        )
      }
    } else {
      // -- In version 3.0.0 this function not used
      callbackUI('', -1, 'inferenceSubVolumes() is not dead code?')
    }
  }
}

async function enableProductionMode(textureF16Flag = true) {
  // -- tf.setBackend('cpu')
  tf.setBackend('webgl')
  // -- tf.removeBackend('cpu')
  // -- Calling enableProdMode() method
  await tf.enableProdMode()
  // -- Setting debug mode of the  environment
  tf.env().set('DEBUG', false)
  tf.env().set('WEBGL_FORCE_F16_TEXTURES', textureF16Flag)
  // -- set this flag so that textures are deleted when tensors are disposed.
  tf.env().set('WEBGL_DELETE_TEXTURE_THRESHOLD', -1)
  // -- tf.env().set('WEBGL_PACK', false)
  // -- Put ready after sets above
  await tf.ready()
  // -- Printing output
  console.log('tf env() flags :', tf.env().flags)
  console.log('tf env() features :', tf.env().features)
  console.log('tf env total features: ', Object.keys(tf.env().features).length)
  console.log('tf backend: ', tf.getBackend())
}

async function runInferenceWW(opts, modelEntry, niftiHeader, niftiImage) {
  const statData = []
  statData.startTime = Date.now() // for common webworker/mainthread do not use performance.now()
  callbackUI('Segmentation started', 0)
  const batchSize = opts.batchSize
  const numOfChan = opts.numOfChan
  if (isNaN(batchSize) || batchSize !== 1) {
    const errTxt = 'The batch Size for input shape must be 1'
    callbackUI(errTxt, -1, errTxt)
    return 0
  }
  if (isNaN(numOfChan) || numOfChan !== 1) {
    const errTxt = 'The number of channels for input shape must be 1'
    callbackUI(errTxt, -1, errTxt)
    return 0
  }
  tf.engine().startScope()
  console.log('Batch size: ', batchSize)
  console.log('Num of Channels: ', numOfChan)
  const model = await load_model(opts.rootURL + modelEntry.path)
  await enableProductionMode(true)
  statData.TF_Backend = tf.getBackend()
  const modelObject = model
  let batchInputShape = []
  // free global variable of 16777216 voxel
  // allOutputSlices3DCC1DimArray = []
  // outputSceneRendered = false
  // read input shape from model.json object
  batchInputShape = modelObject.layers[0].batchInputShape
  console.log(' Model batch input shape : ', batchInputShape)
  // -- Verify input shape
  if (batchInputShape.length !== 5) {
    const errTxt = 'The model input shape must be 5D'
    callbackUI(errTxt, -1, errTxt)
    return 0
  }
  let batch_D, batch_H, batch_W
  const slice_width = niftiHeader.dims[1]
  const slice_height = niftiHeader.dims[2]
  const num_of_slices = niftiHeader.dims[3]
  const isChannelLast = await isModelChnlLast(modelObject)
  if (isChannelLast) {
    console.log('Model Channel Last')
    if (isNaN(batchInputShape[4]) || batchInputShape[4] !== 1) {
      const errTxt = 'The number of channels for input shape must be 1'
      callbackUI(errTxt, -1, errTxt)
      return 0
    }
    batch_D = batchInputShape[1]
    batch_H = batchInputShape[2]
    batch_W = batchInputShape[3]
  } else {
    console.log('Model Channel First')
    if (isNaN(batchInputShape[1]) || batchInputShape[1] !== 1) {
      const errTxt = 'The number of channels for input shape must be 1'
      callbackUI(errTxt, -1, errTxt)
      return 0
    }
    batch_D = batchInputShape[2]
    batch_H = batchInputShape[3]
    batch_W = batchInputShape[4]
  }
  // const input_shape = [batchSize, numOfChan, batch_D, batch_H, batch_W]
  // --Check whether the model will make inference at once as FullVolumeModel
  let isModelFullVol
  if (batch_D === 256 && batch_H === 256 && batch_W === 256) {
    isModelFullVol = true
  } else {
    isModelFullVol = false
  }
  statData.isModelFullVol = isModelFullVol
  // Model output number of segmentations
  let slices_3d = await getAllSlicesDataAsTF3D(num_of_slices, niftiHeader, niftiImage)
  const transpose = modelEntry.enableTranspose
  const enableCrop = modelEntry.enableCrop
  if (isModelFullVol) {
    if (enableCrop) {
      // FullVolume with Crop option before inference ..
      // pre-model to mask the volume, can also be null and the cropping will be on the MRI.
      await inferenceFullVolumePhase1(
        model,
        slices_3d,
        num_of_slices,
        slice_height,
        slice_width,
        isModelFullVol,
        modelEntry,
        statData,
        opts,
        niftiHeader,
        niftiImage
      )
    } else {
      // Transpose MRI data to be match pytorch/keras input output
      console.log('Cropping Disabled')

      if (transpose) {
        slices_3d = slices_3d.transpose()
        console.log('Input transposed')
      } else {
        console.log('Transpose NOT Enabled')
      }

      const enableSeqConv = modelEntry.enableSeqConv

      if (enableSeqConv) {
        callbackUI('', -1, 'inferenceFullVolumeSeqCovLayer() is not dead code?')
      } else {
        callbackUI('', -1, 'inferenceFullVolume() is not dead code?')
      }
    }
  }
}

self.addEventListener(
  'message',
  function (event) {
    runInferenceWW(event.data.opts, event.data.modelEntry, event.data.niftiHeader, event.data.niftiImage)
  },
  false
)