tensor-utils.js

import * as tf from '@tensorflow/tfjs'
import { BWLabeler } from './bwlabels.js'

// DEBUG ONLY: set true to force the post-processing noise guard to trip on any
// run, so you can see the "segmentation produced noise" failure path in the UI
// without an actual garbage segmentation. Leave false for normal use.
const FORCE_SEGMENTATION_NOISE = false;

export async function cropAndGetCorner(tensor3d, mask_3d, userPadding) {
  // Find bounding box
  const [row_min, row_max, col_min, col_max, depth_min, depth_max] = await firstLastNonZero3D(mask_3d);

  // Calculate dimensions
  const height = row_max - row_min + 1;
  const width = col_max - col_min + 1;
  const depth = depth_max - depth_min + 1;

  // Adjust starting corner based on padding, ensuring we don't exceed 256 or go negative
  const adjustCorner = (min, max, size, pad) => {
    const startPad = Math.min(min, pad); // how much we can pad towards start
    const endPad = Math.min(255 - max, pad); // how much we can pad towards end
    const newStart = Math.max(0, min - startPad);
    const newEnd = Math.min(255, max + endPad);
    return [newStart, newEnd];
  };

  const [safeRowStart, safeRowEnd] = adjustCorner(row_min, row_max, height, userPadding);
  const [safeColStart, safeColEnd] = adjustCorner(col_min, col_max, width, userPadding);
  const [safeDepthStart, safeDepthEnd] = adjustCorner(depth_min, depth_max, depth, userPadding);

  // Extract cropped brain with safe bounds
  let cropped = tensor3d.slice(
    [safeRowStart, safeColStart, safeDepthStart],
    [safeRowEnd - safeRowStart + 1, safeColEnd - safeColStart + 1, safeDepthEnd - safeDepthStart + 1]
  );

  // Ensure dimensions are even for stride-2 alignment (important for SpatialAE models)
  // Use zero-padding at the END, matching run_inference.py logic
  const shape = cropped.shape;
  const padRow = shape[0] % 2;
  const padCol = shape[1] % 2;
  const padDepth = shape[2] % 2;

  if (padRow || padCol || padDepth) {
    cropped = cropped.pad([
      [0, padRow],
      [0, padCol],
      [0, padDepth]
    ]);
    console.log(`Padded to even dims: [${shape}] -> [${cropped.shape}]`);
  } else {
    console.log(`Crop dimensions (already even): [${shape}]`);
  }

  return { cropped, corner: [safeRowStart, safeColStart, safeDepthStart], padding: [padRow, padCol, padDepth] };
}




export async function restoreToOriginalSize(tensor3d, corner, targetShape, shift = [0, 0, 0]) {
  const [row_min, col_min, depth_min] = corner;
  const [targetHeight, targetWidth, targetDepth] = targetShape;
  const [height, width, depth] = tensor3d.shape;
  const [sRow, sCol, sDepth] = shift || [0, 0, 0];

  const padRow = Math.max(0, row_min + sRow);
  const padCol = Math.max(0, col_min + sCol);
  const padDepth = Math.max(0, depth_min + sDepth);

  const paddings = [
    [padRow, Math.max(0, targetHeight - height - padRow)],
    [padCol, Math.max(0, targetWidth - width - padCol)],
    [padDepth, Math.max(0, targetDepth - depth - padDepth)]
  ];

  const padded = tensor3d.pad(paddings);
  // Ensure we enforce the target shape (truncate if padding exceeded dimensions)
  if (padded.shape[0] > targetHeight || padded.shape[1] > targetWidth || padded.shape[2] > targetDepth) {
    const sliced = padded.slice([0, 0, 0], [targetHeight, targetWidth, targetDepth]);
    padded.dispose();
    return sliced;
  }
  return padded;
}


export async function addZeroPaddingTo3dTensor(tensor3d, rowPadArr = [1, 1], colPadArr = [1, 1], depthPadArr = [1, 1]) {
  if (tensor3d.rank !== 3) {
    throw new Error('Tensor must be 3D');
  }

  const [height, width, depth] = tensor3d.shape;

  const adjustPadding = (size, padding) => {
    const totalPad = padding[0] + padding[1];
    if (size + totalPad > 256) {
      const scale = (256 - size) / totalPad;
      return [Math.floor(padding[0] * scale), Math.floor(padding[1] * scale)];
    }
    return padding;
  };

  return tensor3d.pad([
    adjustPadding(height, rowPadArr),
    adjustPadding(width, colPadArr),
    adjustPadding(depth, depthPadArr)
  ]);
}

// export async function addZeroPaddingTo3dTensor(tensor3d, rowPadArr = [1, 1], colPadArr = [1, 1], depthPadArr = [1, 1]) {
//   if (tensor3d.rank !== 3) {
//     throw new Error('Tensor must be 3D')
//   }
//   return tensor3d.pad([rowPadArr, colPadArr, depthPadArr])
// }

export async function applyMriThreshold(tensor, percentage) {
  // Perform asynchronous operations outside of tf.tidy
  const maxTensor = tensor.max()
  const thresholdTensor = maxTensor.mul(percentage)
  const threshold = await thresholdTensor.data() // Extracts the threshold value

  // Dispose tensors not needed anymore
  maxTensor.dispose()
  thresholdTensor.dispose()

  // Use tf.tidy for synchronous operations
  return tf.tidy(() => {
    const dataForProcessing = tensor.clone()

    // Thresholding (assuming background has very low values compared to the head)
    const mask = dataForProcessing.greater(threshold[0])
    // -- const denoisedMriData = dataForProcessing.mul(mask)

    // No need to  manually dispose dataForProcessing and mask, as tf.tidy() will dispose them auto.
    return mask
  })

  // -- return denoisedMriData
}

export async function binarizeVolumeDataTensor(volumeDataTensor) {
  const alpha = 0
  // element-wise: (x > 0 ? 1 : alpha * x );  e.g. Tenosr [0, 0.9, 0.8, -3] => Tensor [0, 1, 1, 0]
  return volumeDataTensor.step(alpha)
}

function linearQuantileFromHistogram(histogram, count, quantile) {
  const rank = (count - 1) * quantile
  const lowerRank = Math.floor(rank)
  const upperRank = Math.ceil(rank)

  let cumulative = 0
  let lowerValue = 0
  let upperValue = 0
  let lowerFound = false

  for (let value = 0; value < histogram.length; value++) {
    cumulative += histogram[value]
    if (!lowerFound && cumulative > lowerRank) {
      lowerValue = value
      lowerFound = true
    }
    if (cumulative > upperRank) {
      upperValue = value
      break
    }
  }

  return lowerValue + (upperValue - lowerValue) * (rank - lowerRank)
}

async function calculateQuantiles(tensor, lowerQuantile = 0.02, upperQuantile = 0.98) {
  // Download once to CPU. This avoids tf.gather on a 16M-voxel tensor, which
  // causes avoidable GPU memory pressure immediately before inference.
  const flatTensor = tensor.flatten()
  const totalSize = flatTensor.shape[0]

  // Download tensor data to CPU (TypedArray, not JS Array - much faster)
  const flatData = await flatTensor.data()
  flatTensor.dispose()

  // Brainchomp's bundled examples are uint8. For uint8-like data an exact
  // histogram is both faster and more reproducible than the old random sample,
  // and reproduces NumPy's default linearly interpolated quantile.
  const histogram = new Uint32Array(256)
  let isUint8 = true
  for (let i = 0; i < totalSize; i++) {
    const value = flatData[i]
    if (!Number.isFinite(value)) {
      throw new Error('Cannot percentile-normalize a volume containing NaN or Infinity')
    }
    if (value < 0 || value > 255 || value !== Math.trunc(value)) {
      isUint8 = false
    } else {
      histogram[value]++
    }
  }

  if (isUint8) {
    return {
      qmin: linearQuantileFromHistogram(histogram, totalSize, lowerQuantile),
      qmax: linearQuantileFromHistogram(histogram, totalSize, upperQuantile)
    }
  }

  // Uploaded float volumes can be much larger than is practical to sort in a
  // browser. Use a deterministic, evenly spaced sample instead of Math.random
  // so repeated runs and both browser backends receive the same scaling.
  const sampleSize = Math.min(100000, totalSize)
  const sample = new Float32Array(sampleSize)
  const scale = sampleSize > 1 ? (totalSize - 1) / (sampleSize - 1) : 0
  for (let i = 0; i < sampleSize; i++) {
    const value = flatData[Math.round(i * scale)]
    if (!Number.isFinite(value)) {
      throw new Error('Cannot percentile-normalize a volume containing NaN or Infinity')
    }
    sample[i] = value
  }
  sample.sort()

  const interpolate = (quantile) => {
    const rank = (sampleSize - 1) * quantile
    const lower = Math.floor(rank)
    const upper = Math.ceil(rank)
    return sample[lower] + (sample[upper] - sample[lower]) * (rank - lower)
  }

  const qminValue = interpolate(lowerQuantile)
  const qmaxValue = interpolate(upperQuantile)

  return { qmin: qminValue, qmax: qmaxValue }
}

export async function convByOutputChannelAndInputSlicing(input, filter, biases, stride, pad, dilationRate, sliceSize) {
  const inChannels = input.shape[4]
  const outChannels = filter.shape[4]

  // Create an empty array to hold the output channels
  let outputChannels = null

  // Slice the input tensor and process one output channel at a time
  for (let channel = 0; channel < outChannels; channel++) {
    const numSlices = Math.ceil(inChannels / sliceSize)
    let outputChannel = null

    for (let i = 0; i < numSlices; i++) {
      const startChannel = i * sliceSize
      const endChannel = Math.min((i + 1) * sliceSize, inChannels)

      // Only proceed if there are channels to process
      if (startChannel < inChannels) {
        const resultSlice = tf.tidy(() => {
          const inputSlice = input.slice([0, 0, 0, 0, startChannel], [-1, -1, -1, -1, endChannel - startChannel])
          const filterSlice = filter.slice([0, 0, 0, startChannel, channel], [-1, -1, -1, endChannel - startChannel, 1])
          // Perform the convolution for the current slice and output channel
          return tf.conv3d(inputSlice, filterSlice, stride, pad, 'NDHWC', dilationRate)
        })

        if (outputChannel === null) {
          outputChannel = resultSlice
        } else {
          const updatedOutputChannel = outputChannel.add(resultSlice)
          outputChannel.dispose()
          resultSlice.dispose()
          outputChannel = updatedOutputChannel
        }
      }
    }

    // --- Start of Fix ---
    // This variable will hold the final result for this channel, after bias is (optionally) added.
    let biasedOutputChannel;

    // Check if the biases tensor was actually provided for this layer.
    if (biases) {
      // If biases exist, slice the correct one for this channel and add it.
      const biasesSlice = biases.slice([channel], [1]);
      biasedOutputChannel = outputChannel.add(biasesSlice);
      outputChannel.dispose();
      biasesSlice.dispose();
    } else {
      // If no biases exist, the result is simply the accumulated convolution output.
      biasedOutputChannel = outputChannel;
    }
    // --- End of Fix ---

    // Accumulate the channel to the output array
    if (outputChannels == null) {
      outputChannels = biasedOutputChannel
    } else {
      const updatedOutputChannels = await tf.concat([outputChannels, biasedOutputChannel], 4)
      biasedOutputChannel.dispose()
      outputChannels.dispose()
      outputChannels = updatedOutputChannels
    }
  }

  return outputChannels
}

/**
 * Applies instance normalization to a single-channel 3D tensor.
 * Normalizes the tensor over its depth, height, and width.
 * @param {tf.Tensor} x A 5D tensor of shape [1, D, H, W, 1].
 * @param {number} epsilon A small float added to variance to avoid dividing by zero.
 * @returns {tf.Tensor} The normalized tensor.
 */
function instanceNorm(x, epsilon = 1e-5) {
  return tf.tidy(() => {
    // Axes [1, 2, 3] correspond to Depth, Height, Width.
    const { mean, variance } = tf.moments(x, [1, 2, 3], true);
    const invStd = tf.rsqrt(variance.add(epsilon));
    return x.sub(mean).mul(invStd);
  });
}


export async function gn_convByOutputChannelAndInputSlicing(input, filter, biases, stride, pad, dilationRate, sliceSize) {

  //    const finalResult = tf.tidy(() => {
  const inChannels = input.shape[4];
  const outChannels = filter.shape[4];

  let outputChannels = null;

  for (let channel = 0; channel < outChannels; channel++) {
    // This section computes a single channel's convolution.
    const numSlices = Math.ceil(inChannels / sliceSize);
    let outputChannel = null;
    for (let i = 0; i < numSlices; i++) {
      const startChannel = i * sliceSize;
      const endChannel = Math.min((i + 1) * sliceSize, inChannels);
      if (startChannel < inChannels) {
        const resultSlice = tf.tidy(() => {
          const inputSlice = input.slice([0, 0, 0, 0, startChannel], [-1, -1, -1, -1, endChannel - startChannel]);
          const filterSlice = filter.slice([0, 0, 0, startChannel, channel], [-1, -1, -1, endChannel - startChannel, 1]);
          return tf.conv3d(inputSlice, filterSlice, stride, pad, 'NDHWC', dilationRate);
        });
        if (outputChannel === null) {
          outputChannel = resultSlice;
        } else {
          const updatedOutputChannel = outputChannel.add(resultSlice);
          outputChannel.dispose();
          resultSlice.dispose();
          outputChannel = updatedOutputChannel;
        }
      }
    }

    let biasedOutputChannel;
    if (biases) {
      const biasesSlice = biases.slice([channel], [1]);
      biasedOutputChannel = outputChannel.add(biasesSlice);
      outputChannel.dispose();
      biasesSlice.dispose();
    } else {
      biasedOutputChannel = outputChannel;
    }

    // --- KEY CHANGE ---
    // Apply instance normalization to the resulting channel.
    // Apply normalization
    const normalizedChannel = instanceNorm(biasedOutputChannel);
    biasedOutputChannel.dispose();

    // Incremental concatenation
    if (outputChannels === null) {
      outputChannels = normalizedChannel;
    } else {
      const updatedOutputChannels = await tf.concat([outputChannels, normalizedChannel], 4);
      normalizedChannel.dispose();
      outputChannels.dispose();
      outputChannels = updatedOutputChannels;
    }
  }

  return outputChannels;
}// );
//     return finalResult;
// }

// ===========================================================================
// CHANNEL-LIST (sequential) convolution path
// ---------------------------------------------------------------------------
// The functions below carry a layer's activation as a JavaScript ARRAY of
// single-channel tensors ([1, D, H, W, 1] each) instead of one packed
// [1, D, H, W, C] tensor. For the large gridding-free MeshNet models the full
// multi-channel activation (~256x204x204x24) would be materialized by tfjs as
// an unpacked WebGL2 texture larger than the 8192 limit and throw. A single
// channel is only ~3257^2 packed, comfortably under the limit -- so as long as
// no tensor ever holds more than one channel, every texture fits.
//
// These replace the old convByOutputChannelAndInputSlicing / gn_ variants,
// whose final `tf.concat([...], 4)` rebuilt the oversized full tensor.
// ===========================================================================

/**
 * Compute ONE output channel of a Conv3D as a single-channel [1, D, H, W, 1]
 * tensor, reading inputs from a channel-list (array of single-channel tensors).
 * Input channels are gathered in small groups of `sliceSize` so the only
 * multi-channel tensor ever built is a tiny [1, D, H, W, sliceSize] temporary.
 *
 * @param {tf.Tensor[]} inputList Array of [1, D, H, W, 1] input channels.
 * @param {tf.Tensor} filter Conv3D kernel [kD, kH, kW, inC, outC].
 * @param {tf.Tensor|null} biases Bias vector [outC] or null.
 * @param {number} channel Index of the output channel to compute.
 * @returns {tf.Tensor} A kept [1, D, H, W, 1] tensor (caller disposes).
 */
function convOneOutputChannel(inputList, filter, biases, channel, stride, pad, dilationRate, sliceSize) {
  const inChannels = inputList.length;
  return tf.tidy(() => {
    let acc = null;
    const numSlices = Math.ceil(inChannels / sliceSize);
    for (let i = 0; i < numSlices; i++) {
      const start = i * sliceSize;
      const end = Math.min((i + 1) * sliceSize, inChannels);
      const group = end - start;
      // Gather this small group of input channels. group===1 avoids a needless
      // concat (and never materializes anything wider than `sliceSize`).
      const inputSlice = group === 1
        ? inputList[start]
        : tf.concat(inputList.slice(start, end), 4);
      const filterSlice = filter.slice([0, 0, 0, start, channel], [-1, -1, -1, group, 1]);
      const part = tf.conv3d(inputSlice, filterSlice, stride, pad, 'NDHWC', dilationRate);
      acc = acc === null ? part : acc.add(part);
    }
    if (biases) {
      acc = acc.add(biases.slice([channel], [1]));
    }
    return acc;
  });
}

/**
 * Channel-list Conv3D: returns an ARRAY of single-channel [1, D, H, W, 1]
 * tensors (one per output channel) instead of a concatenated full tensor.
 * Optionally applies per-channel normalization (the affine-GroupNorm export
 * decomposes GroupNorm into per-channel instance-norm + a diagonal 1x1
 * "affine_*" conv, so per-channel normalization here is mathematically exact).
 * Uses the CENTERED-variance LayerNormInPlace -- the same normalization the
 * dense WebGL path uses, which is required for fp16-texture stability (the
 * one-pass E[x^2]-E[x]^2 form cancels catastrophically in fp16 and produces
 * pure-noise segmentations).
 *
 * Caller owns the returned array and must dispose it (tf.dispose(list)).
 *
 * @param {tf.Tensor[]} inputList Array of [1, D, H, W, 1] input channels.
 * @param {tf.Tensor} filter Conv3D kernel [kD, kH, kW, inC, outC].
 * @param {tf.Tensor|null} biases Bias vector [outC] or null.
 * @param {boolean} applyNorm Apply per-channel GroupNorm (instance norm) to each output channel.
 * @returns {tf.Tensor[]} Array of `outC` single-channel tensors.
 */
export function convChannelList(inputList, filter, biases, stride, pad, dilationRate, sliceSize, applyNorm = false) {
  const outChannels = filter.shape[4];
  const outputList = [];
  for (let c = 0; c < outChannels; c++) {
    let ch = convOneOutputChannel(inputList, filter, biases, c, stride, pad, dilationRate, sliceSize);
    if (applyNorm) {
      const normed = LayerNormInPlace(ch); // centered-variance, fp16-safe, per-channel
      ch.dispose();
      ch = normed;
    }
    outputList.push(ch);
  }
  return outputList;
}

/**
 * Channel-list transposed (strided up-sampling) Conv3D, for SpatialAE-style
 * MeshNet variants (e.g. Tissue GWM / model_sae16ch3). Returns an ARRAY of
 * single-channel [1, D', H', W', 1] tensors (one per output channel) without
 * ever materializing a full multi-channel activation -- so the up-sampled
 * 256^3 result stays under the WebGL2 8192 texture limit.
 *
 * The tfjs/Keras Conv3DTranspose kernel layout is [kD, kH, kW, outC, inC]
 * (out/in swapped vs Conv3D). Each output channel is the sum over input
 * channels of conv3dTranspose(inChannel, kernel[..., o, c]).
 *
 * @param {tf.Tensor[]} inputList Array of [1, D, H, W, 1] input channels.
 * @param {tf.Tensor} filter Conv3DTranspose kernel [kD, kH, kW, outC, inC].
 * @param {tf.Tensor|null} biases Bias vector [outC] or null.
 * @param {number[]} outSpatialShape Output [D', H', W'] (from layer.computeOutputShape).
 * @returns {tf.Tensor[]} Array of `outC` single-channel tensors.
 */
export function convTransposeChannelList(inputList, filter, biases, outSpatialShape, stride, pad) {
  const outChannels = filter.shape[3];
  const inChannels = filter.shape[4];
  const outShape5d = [1, outSpatialShape[0], outSpatialShape[1], outSpatialShape[2], 1];
  const outputList = [];
  for (let o = 0; o < outChannels; o++) {
    const ch = tf.tidy(() => {
      let acc = null;
      for (let c = 0; c < inChannels; c++) {
        const kSlice = filter.slice([0, 0, 0, o, c], [-1, -1, -1, 1, 1]); // [kD,kH,kW,1,1]
        const part = tf.conv3dTranspose(inputList[c], kSlice, outShape5d, stride, pad);
        acc = acc === null ? part : acc.add(part);
      }
      if (biases) {
        acc = acc.add(biases.slice([o], [1]));
      }
      return acc;
    });
    outputList.push(ch);
  }
  return outputList;
}

/**
 * Final classifier + argmax over a channel-list input, without ever building
 * the full [1, D, H, W, numClasses] logits tensor. Computes each class logit as
 * a single-channel tensor and folds it into a running (max-logit, argmax-index)
 * pair -- the same incremental-argmax trick as SequentialConvLayer, but reading
 * from a channel-list instead of a packed tensor.
 *
 * @param {tf.Tensor[]} inputList Array of [1, D, H, W, 1] activations (backbone output).
 * @param {tf.Tensor} weights Final Conv3D kernel [kD, kH, kW, inC, numClasses].
 * @param {tf.Tensor|null} biases Bias vector [numClasses] or null.
 * @param {function} [callbackUI] Optional progress callback (msg, frac).
 * @returns {Promise<tf.Tensor>} Squeezed [D, H, W] argmax label volume (float32).
 */
export async function seqConvArgMaxChannelList(inputList, weights, biases, stride, pad, dilationRate, callbackUI, isWebWorker = true) {
  const numClasses = weights.shape[4];
  const sliceSize = 3;
  // The running max/argmax are kept at RANK 3 ([D, H, W]). tf.where (select) on
  // the WebGL backend only supports up to rank 4 -- Firefox throws "Where for
  // rank 5 is not yet supported" on the [1, D, H, W, 1] form -- so we reshape
  // each single-channel logit down to [D, H, W] before the compare/select.
  let outB = null; // running max logit  [D, H, W]
  let outC = null; // running argmax idx [D, H, W]
  let spatialShape = null;

  for (let k = 0; k < numClasses; k++) {
    const logit5d = convOneOutputChannel(inputList, weights, biases, k, stride, pad, dilationRate, sliceSize);
    if (spatialShape === null) {
      // [1, D, H, W, 1] -> [D, H, W]
      spatialShape = [logit5d.shape[1], logit5d.shape[2], logit5d.shape[3]];
    }
    const logit = tf.tidy(() => logit5d.reshape(spatialShape));
    logit5d.dispose();

    if (outB === null) {
      outB = logit;
      outC = tf.zerosLike(logit);
    } else {
      const [newB, newC] = tf.tidy(() => {
        const greater = tf.greater(logit, outB);
        return [tf.where(greater, logit, outB), tf.where(greater, tf.fill(outC.shape, k), outC)];
      });
      outB.dispose();
      outC.dispose();
      logit.dispose();
      outB = newB;
      outC = newC;
    }
    if (callbackUI) callbackUI(`Final layer class ${k + 1}/${numClasses}`, (k + 1) / numClasses);
    // Yield to the event loop periodically so the UI can update / avoid GPU
    // watchdog timeouts on the long classifier loop (only on main thread).
    if (!isWebWorker && (k % 8 === 0)) {
      await new Promise(resolve => setTimeout(resolve, 0));
    }
  }

  outB.dispose();
  return outC; // already [D, H, W]
}

/**
 * Split a packed [1, D, H, W, C] tensor into a channel-list of C single-channel
 * [1, D, H, W, 1] tensors. For C===1 the input is returned as a one-element list
 * WITHOUT copying (the caller then owns disposal via the list).
 *
 * @param {tf.Tensor} tensor Packed [1, D, H, W, C] tensor.
 * @returns {tf.Tensor[]} Array of C single-channel tensors.
 */
export function splitToChannelList(tensor) {
  const C = tensor.shape[4];
  if (C === 1) return [tensor];
  const list = [];
  for (let c = 0; c < C; c++) {
    list.push(tensor.slice([0, 0, 0, 0, c], [-1, -1, -1, -1, 1]));
  }
  return list;
}

/**
 * Applies instance normalization to a tensor and disposes the input tensor
 * to simulate an in-place operation, minimizing peak memory usage.
 * Assumes input shape [1, D, H, W, C]. Normalization is per-channel.
 *
 * @param {tf.Tensor} x The input tensor to normalize.
 * @param {number} [epsilon=1e-5] A small float added to variance to avoid division by zero.
 * @returns {tf.Tensor} The new, normalized tensor.
 */
export function LayerNormInPlace(x, epsilon = 1e-5) {
  // Per-channel z-score over the spatial axes (NDHWC: reduce D,H,W; keep N,C).
  //
  // The straightforward `tf.moments(x, [1,2,3])` is expensive on WebGL: because
  // the channel axis is innermost and is KEPT, each of its two internal
  // reductions (mean, then mean-of-squared-deviations) transposes the full
  // ~256M-element volume, and it also materializes a full (x - mean) volume.
  // That cost ~965 ms/layer x 13 layers in profiling.
  //
  // Here we transpose ONCE to channel-first, take both reductions over the now
  // innermost (contiguous) spatial block -- so they need no further transpose --
  // and compute variance in a single pass as E[x^2] - E[x]^2. The per-channel
  // [C] stats then broadcast directly over the innermost axis of the original
  // NDHWC tensor, so the final normalize needs no transpose either. Result is
  // numerically equivalent (variance clamped >= 0 to absorb fp round-off).
  // IMPORTANT: use the CENTERED two-pass variance (mean of (x-mean)^2), NOT the
  // one-pass E[x^2]-E[x]^2. The one-pass form is mathematically identical but in
  // fp16 textures it cancels catastrophically (E[x^2] ~ E[x]^2, both large), which
  // produced pure-noise segmentations (~1.5M connected components instead of ~2.5k).
  return tf.tidy(() => {
    const rank = x.shape.length;            // expected 5: [1, D, H, W, C]
    const C = x.shape[rank - 1];
    const N = x.shape[1] * x.shape[2] * x.shape[3]; // D * H * W

    // One transpose to channel-first, then both reductions run over the now
    // innermost (contiguous) spatial axis -- no further transpose. tf.moments
    // would transpose the full volume twice (once per internal reduction).
    const flatCN = x.transpose([0, 4, 1, 2, 3]).reshape([C, N]); // [C, N]
    const mean = flatCN.mean(1);                       // E[x]            [C]
    const centered = flatCN.sub(mean.reshape([C, 1]));  // x - mean        [C, N]
    const variance = centered.square().mean(1);        // E[(x-mean)^2]   [C]
    const invStd = tf.rsqrt(tf.add(variance, epsilon)); // [C]

    const meanB = mean.reshape([1, 1, 1, 1, C]);
    const invStdB = invStd.reshape([1, 1, 1, 1, C]);
    return x.sub(meanB).mul(invStdB);
  });
}

export async function draw3dObjBoundingVolume(unstackOutVolumeTensor, opts, modelEntry, callbackImg) {
  const allOutputSlices3DCC = []

  // dataSync() using to flatten array. Takes around 1.5 s
  for (let sliceTensorIdx = 0; sliceTensorIdx < unstackOutVolumeTensor.length; sliceTensorIdx++) {
    allOutputSlices3DCC[sliceTensorIdx] = Array.from(unstackOutVolumeTensor[sliceTensorIdx].dataSync())
  }

  // Use this conversion to download output slices as nii file. Takes around 30 ms
  // does not use `push` to avoid stack overflows. In future: consider .set() with typed arrays
  const allOutputSlices3DCC1DimArray = new Array(allOutputSlices3DCC[0].length * allOutputSlices3DCC.length)
  let index = 0
  for (let sliceIdx = 0; sliceIdx < allOutputSlices3DCC.length; sliceIdx++) {
    for (let i = 0; i < allOutputSlices3DCC[sliceIdx].length; i++) {
      allOutputSlices3DCC1DimArray[index++] = allOutputSlices3DCC[sliceIdx][i]
    }
  }
  console.log('Done with allOutputSlices3DCC1DimArray ')
  const brainMaskTensor1d = await binarizeVolumeDataTensor(tf.tensor1d(allOutputSlices3DCC1DimArray))
  const brainOut = Array.from(brainMaskTensor1d.dataSync())
  callbackImg(brainOut, opts, modelEntry)
}
// return first and last non-zero voxel in row (dim = 0), column (1) or slice (2) dimension
async function firstLastNonZero(tensor3D, dim = 0) {
  let mxs = []
  if (dim === 0) {
    mxs = await tensor3D.max(2).max(1).arraySync()
  } else if (dim === 1) {
    mxs = await tensor3D.max(2).max(0).arraySync()
  } else {
    mxs = await tensor3D.max(1).max(0).arraySync()
  }
  let mn = mxs.length
  let mx = 0
  for (let i = 0; i < mxs.length; i++) {
    if (mxs[i] > 0) {
      mn = i
      break
    }
  }
  for (let i = mxs.length - 1; i >= 0; i--) {
    if (mxs[i] > 0) {
      mx = i
      break
    }
  }
  return [mn, mx]
}

export async function firstLastNonZero3D(tensor3D) {
  const [row_min, row_max] = await firstLastNonZero(tensor3D, 0)
  const [col_min, col_max] = await firstLastNonZero(tensor3D, 1)
  const [depth_min, depth_max] = await firstLastNonZero(tensor3D, 2)
  console.log('row min and max  :', row_min, row_max)
  console.log('col min and max  :', col_min, col_max)
  console.log('depth min and max  :', depth_min, depth_max)
  return [row_min, row_max, col_min, col_max, depth_min, depth_max]
}

/*
//simpler function, but x4 slower
export async function firstLastNonZero3D(tensor3D) {
  const coords = await tf.whereAsync(tensor3D)
  const row_min = coords.min(0).arraySync()[0]
  const row_max = coords.max(0).arraySync()[0]
  const col_min = coords.min(0).arraySync()[1]
  const col_max = coords.max(0).arraySync()[1]
  const depth_min = coords.min(0).arraySync()[2]
  const depth_max = coords.max(0).arraySync()[2]
  coords.dispose()
  return [row_min, row_max, col_min, col_max, depth_min, depth_max]
}
*/

export async function generateBrainMask(
  unstackOutVolumeTensor,
  num_of_slices,
  slice_height,
  slice_width,
  modelEntry,
  opts,
  callbackUI,
  callbackImg,
  isFinalImage = true
) {
  if (unstackOutVolumeTensor[0].dtype !== 'int32') {
    callbackUI('', -1, 'generateBrainMask assumes int32')
  }
  if (modelEntry.preModelPostProcess) {
    callbackUI('', -1, 'generateBrainMask assumes BWLabeler instead of preModelPostProcess')
  }
  const numSlices = unstackOutVolumeTensor.length
  const numPixels2D = unstackOutVolumeTensor[0].size
  const numVox3D = numSlices * numPixels2D
  // preallocate to reduce heap usage
  const brainOut = new Int32Array(numVox3D)
  let offset = 0
  for (let i = 0; i < numSlices; i++) {
    brainOut.set(unstackOutVolumeTensor[i].dataSync(), offset)
    offset += numPixels2D
  }
  for (let i = 0; i < numVox3D; i++) {
    brainOut[i] = brainOut[i] !== 0 ? 1 : 0
  }
  if (isFinalImage || opts.showPhase1Output) {
    // all done
    callbackImg(brainOut, opts, modelEntry)
    callbackUI('Segmentation finished', 0)
  }
  return tf.tensor(brainOut, [num_of_slices, slice_height, slice_width])
}

export async function generateOutputSlicesV2(
  img,
  OutVolumeTensorShape,
  OutVolumeTensorType,
  num_of_slices,
  numSegClasses,
  slice_height,
  slice_width,
  modelEntry,
  opts,
  niftiImage
) {
  // Convert all slices into 1 Dim array
  if (opts.isPostProcessEnable) {
    const BWInstance = new BWLabeler()
    const dim = new Uint32Array(OutVolumeTensorShape)
    const conn = 6 // Example connectivity
    const binarize = true
    const onlyLargestClusterPerClass = true
    const [_labelCount, labeledImage] = BWInstance.bwlabel(img, dim, conn, binarize, onlyLargestClusterPerClass)
    for (let i = 0; i < img.length; i++) {
      img[i] *= labeledImage[i]
    }
  } // if isPostProcessEnable
  const typedArrayConstructor = {
    float32: Float32Array,
    int32: Int32Array
    // Add other cases as needed for different dtypes
  }[OutVolumeTensorType]
  // Create a new TypedArray from img with the same type as outLabelVolume
  const allOutputSlices3DCC1DimArray = new Uint8Array(img)
  switch (modelEntry.type) {
    case 'Brain_Masking': {
      const brainMask = new Uint8Array(allOutputSlices3DCC1DimArray.length)
      for (let i = 0; i < allOutputSlices3DCC1DimArray.length; i++) {
        brainMask[i] = allOutputSlices3DCC1DimArray[i] !== 0 ? 1 : 0
      }
      return brainMask
    }
    case 'Brain_Extraction': {
      if (modelEntry.returnMaskForExtraction) {
        const brainMask = new Uint8Array(allOutputSlices3DCC1DimArray.length)
        for (let i = 0; i < allOutputSlices3DCC1DimArray.length; i++) {
          brainMask[i] = allOutputSlices3DCC1DimArray[i] !== 0 ? 1 : 0
        }
        return brainMask
      }
      const maskedData = new niftiImage.constructor(allOutputSlices3DCC1DimArray.length)
      for (let i = 0; i < allOutputSlices3DCC1DimArray.length; i++) {
        // Create the mask - 1 where the value is non-zero, 0 where it is zero.
        const maskValue = allOutputSlices3DCC1DimArray[i] !== 0 ? 1 : 0
        // Apply the mask to the data - multiply by the mask value.
        maskedData[i] = niftiImage[i] * maskValue
      }
      return maskedData
    }
  }
  return img
}

export async function getAllSlicesDataAsTF3D(num_of_slices, niftiHeader, niftiImage) {
  // Get nifti dimensions
  const cols = niftiHeader.dims[1] // Slice width
  const rows = niftiHeader.dims[2] // Slice height
  let typedData
  if (niftiHeader.datatypeCode === 2) {
    // enum from nvimage/utils DT_UINT8 = 2
    typedData = new Uint8Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 4) {
    // DT_INT16 = 4
    typedData = new Int16Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 8) {
    // DT_INT32 = 8
    typedData = new Int32Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 16) {
    // DT_FLOAT32 = 16
    typedData = new Float32Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 64) {
    // DT_FLOAT64 = 64
    typedData = new Float64Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 256) {
    // DT_INT8 = 256
    typedData = new Int8Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 512) {
    // DT_UINT16 = 512
    typedData = new Uint16Array(niftiImage)
  } else if (niftiHeader.datatypeCode === 768) {
    // DT_UINT32 = 768
    typedData = new Uint32Array(niftiImage)
  } else {
    return
  }
  const allSlices_2D = []
  let offset3D = 0
  // Draw pixels
  for (let slice = 0; slice < num_of_slices; slice++) {
    const slice = new Array(rows * cols)
    let offset2D = 0
    for (let row = 0; row < rows; row++) {
      for (let col = 0; col < cols; col++) {
        const value = typedData[offset3D++]
        // Create 1Dim Array of pixel value, this 1 dim represents one channel
        slice[offset2D++] = value & 0xff
      }
    }
    allSlices_2D.push(tf.tensor(slice, [rows, cols])) // slice_height, slice_width
  }
  const allSlices_3D = tf.stack(allSlices_2D)
  tf.dispose(allSlices_2D)
  return allSlices_3D
}

export async function getModelNumLayers(modelObj) {
  return modelObj.layers.length
}

export async function getModelNumParameters(modelObj) {
  let numParameters = 0
  for (let layerIdx = 0; layerIdx < modelObj.layers.length; layerIdx++) {
    numParameters += modelObj.layers[layerIdx].countParams()
  }
  return numParameters
}

export async function isModelChnlLast(modelObj) {
  for (let layerIdx = 0; layerIdx < modelObj.layers.length; layerIdx++) {
    if (modelObj.layersByDepth[layerIdx][0].dataFormat) {
      return modelObj.layersByDepth[layerIdx][0].dataFormat === 'channelsLast'
    }
  }
}

export async function load_model(modelUrl) {
  return await tf.loadLayersModel(modelUrl)
}

export async function minMaxNormalizeVolumeData(volumeData) {
  // Normalize the data to the range 0 - 1 using min-max scaling
  const volumeData_Max = volumeData.max()
  const volumeData_Min = volumeData.min()
  const normalizedSlices_3d = await volumeData.sub(volumeData_Min).div(volumeData_Max.sub(volumeData_Min))
  return normalizedSlices_3d
}

function processTensorInChunks(inputTensor, filterWeights, chunkSize) {
  // Assuming inputTensor's shape: [batch, depth, height, width, inChannels]
  // and filterWeights's shape: [filterDepth, filterHeight, filterWidth, inChannels, outChannels]
  const stride = 1
  const pad = 0
  const dilationRate = 1
  const inChannels = inputTensor.shape[4]
  const numSlices = Math.ceil(inChannels / chunkSize)

  let accumulatedResult = null
  for (let i = 0; i < numSlices; i++) {
    const startChannel = i * chunkSize
    const endChannel = Math.min((i + 1) * chunkSize, inChannels)
    const channels = endChannel - startChannel

    const inputSlice = tf.tidy(() => {
      // Slice the input tensor to get the current chunk
      return inputTensor.slice([0, 0, 0, 0, startChannel], [-1, -1, -1, -1, channels])
    })
    const filterSlice = tf.tidy(() => {
      // Slice the filter weights to match the input tensor's current chunk
      return filterWeights.slice([0, 0, 0, startChannel, 0], [-1, -1, -1, channels, -1])
    })

    const resultSlice = tf.conv3d(inputSlice, filterSlice, stride, pad, 'NDHWC', dilationRate)
    // Clean up the slices to free memory
    inputSlice.dispose()
    filterSlice.dispose()

    // Squeeze the result slice to remove dimensions of size 1
    const squeezedResultSlice = tf.squeeze(resultSlice)
    resultSlice.dispose() // Dispose of the original resultSlice after squeezing

    if (accumulatedResult === null) {
      accumulatedResult = squeezedResultSlice
    } else {
      // Accumulate the result by adding the new result slice to it
      const newAccumulatedResult = accumulatedResult.add(squeezedResultSlice)

      // Dispose of the previous accumulatedResult and squeezedResultSlice
      accumulatedResult.dispose()
      // Dispose of squeezedResultSlice only if it wasn't assigned to accumulatedResult
      if (accumulatedResult !== squeezedResultSlice) {
        squeezedResultSlice.dispose()
      }
      // Update accumulatedResult with the new result
      accumulatedResult = newAccumulatedResult
    }

    tf.tidy(() => {
      tf.matMul(tf.zeros([1, 1]), tf.zeros([1, 1]))
    })
  }

  return accumulatedResult
}

export async function quantileNormalizeVolumeData(
  tensor,
  lowerQuantile = 0.02,
  upperQuantile = 0.98,
  denominatorEpsilon = 1e-3
) {
  if (!(lowerQuantile >= 0 && lowerQuantile < upperQuantile && upperQuantile <= 1)) {
    throw new Error(`Invalid normalization percentiles: ${lowerQuantile}, ${upperQuantile}`)
  }

  // Call calculateQuantiles and wait for the result
  const { qmin, qmax } = await calculateQuantiles(tensor, lowerQuantile, upperQuantile)
  console.log(
    `[Normalization] ${lowerQuantile * 100}-${upperQuantile * 100} percentiles: ` +
    `low=${qmin}, high=${qmax}, epsilon=${denominatorEpsilon}, clip=[0,1]`
  )

  // Perform the operation: (tensor - qmin) / (qmax - qmin)
  // Break up chained operations to properly dispose intermediate tensors
  const range = qmax - qmin + denominatorEpsilon
  const shifted = tensor.sub(qmin)
  const scaled = shifted.div(range)
  const resultTensor = scaled.clipByValue(0, 1)
  shifted.dispose() // Dispose intermediate tensor to prevent memory leak
  scaled.dispose()

  // Return the resulting tensor (caller is responsible for disposing input tensor)
  return resultTensor
}

export async function removeZeroPaddingFrom3dTensor(tensor3d, rowPad = 1, colPad = 1, depthPad = 1) {
  if (tensor3d.rank !== 3) {
    throw new Error('Tensor must be 3D')
  }
  const [h, w, d] = tensor3d.shape
  return tensor3d.slice([rowPad, colPad, depthPad], [h - 2 * rowPad, w - 2 * colPad, d - 2 * depthPad])
}

export async function resizeWithZeroPadding(croppedTensor3d, newDepth, newHeight, newWidth, refVoxel, boundVolSizeArr) {
  const row_pad_befor = refVoxel[0]
  const col_pad_befor = refVoxel[1]
  const depth_pad_befor = refVoxel[2]
  // last and lower volume voxel
  const row_max = row_pad_befor + boundVolSizeArr[0] - 1 // size [2, 2, 2] means 2 voxels total in each dim
  const col_max = col_pad_befor + boundVolSizeArr[1] - 1
  const depth_max = depth_pad_befor + boundVolSizeArr[2] - 1

  const row_pad_after = newHeight - row_max - 1 > 0 ? newHeight - row_max - 1 : 0
  const col_pad_after = newWidth - col_max - 1 > 0 ? newWidth - col_max - 1 : 0
  const depth_pad_after = newDepth - depth_max - 1 > 0 ? newDepth - depth_max - 1 : 0

  return croppedTensor3d.pad([
    [row_pad_befor, row_pad_after],
    [col_pad_befor, col_pad_after],
    [depth_pad_befor, depth_pad_after]
  ])
}


export class SequentialConvLayer {
  constructor(model, chunkSize, isChannelLast, callbackUI, isWebWorker = true) {
    this.model = model;
    this.outChannels = model.outputLayers[0].kernel.shape[4];
    this.chunkSize = chunkSize;
    this.isChannelLast = isChannelLast;
    this.callbackUI = callbackUI;
    this.isWebWorker = isWebWorker;
  }

  async apply(inputTensor) {
    const startTime = performance.now();
    const convLayer = this.model.layers[this.model.layers.length - 1];
    const weights = convLayer.getWeights()[0];
    const biases = convLayer.getWeights()[1];
    const outputShape = this.isChannelLast ? inputTensor.shape.slice(1, -1) : inputTensor.shape.slice(2);

    // Initialize output tensors
    let outB = await tf.mul(tf.ones(outputShape), -10000);
    let outC = await tf.zeros(outputShape);

    // Process in larger chunks to reduce overhead
    const CHUNK_SIZE = 3; // Process 8 channels at once
    const chunks = Math.ceil(this.outChannels / CHUNK_SIZE);

    for (let chunk = 0; chunk < chunks; chunk++) {
      const startIdx = chunk * CHUNK_SIZE;
      const endIdx = Math.min((chunk + 1) * CHUNK_SIZE, this.outChannels);

      // Process chunk in single tidy to optimize memory
      const [newOutB, newOutC] = await tf.tidy(() => {
        let currentOutB = outB;
        let currentOutC = outC;

        for (let chIdx = startIdx; chIdx < endIdx; chIdx++) {
          const filterWeights = weights.slice([0, 0, 0, 0, chIdx], [-1, -1, -1, -1, 1]);
          const filterBiases = biases.slice([chIdx], [1]);

          const outA = processTensorInChunks(
            inputTensor,
            filterWeights,
            Math.min(this.chunkSize, this.outChannels)
          ).add(filterBiases);

          const greater = tf.greater(outA, currentOutB);
          currentOutB = tf.where(greater, outA, currentOutB);
          currentOutC = tf.where(greater, tf.fill(currentOutC.shape, chIdx), currentOutC);
        }

        return [currentOutB, currentOutC];
      });

      // Clean up previous tensors
      tf.dispose([outB, outC]);
      outB = newOutB;
      outC = newOutC;

      // Update progress
      this.callbackUI(`Processing chunk ${chunk + 1}/${chunks}`, (chunk + 1) / chunks);

      // Allow UI update and prevent GPU timeout
      if (!this.isWebWorker) {
        await new Promise(resolve => setTimeout(resolve, 0));
      }
    }

    // Clean up and return result
    const result = outC.clone();
    tf.dispose([outB, outC]);

    const endTime = performance.now();
    console.log(`Execution time: ${endTime - startTime} milliseconds`);
    return result;
  }
}



/**
 * Main orchestrator function to process a segmentation volume.
 * OPTIMIZED: This version performs a single GPU->CPU data transfer and then
 * executes all subsequent logic on the CPU, avoiding inefficient data round-trips.
 * It is based on the logic of the proven `generateOutputSlicesV2` function.
 *
 * @param {tf.Tensor} outLabelVolume The final tensor from the model inference.
 * @param {Uint8Array|Float32Array} niftiImage The raw data from the original NIfTI file.
 * @param {object} modelEntry The model configuration object.
 * @param {object} opts The options object, containing `isPostProcessEnable`.
 * @returns {Promise<Uint8Array>} The final processed image data as a Uint8Array.
 */
export async function processSegmentationVolume(outLabelVolume, niftiImage, modelEntry, opts) {
  // --- Step 1: Single Data Transfer from GPU to CPU ---
  console.log('Downloading segmentation data from GPU to CPU...');
  const segmentationData = await outLabelVolume.data(); // This returns a TypedArray (e.g., Int32Array)

  const Vshape = outLabelVolume.shape;
  console.log('Data download complete. Starting CPU processing.');

  // --- Step 2: Apply Post-Processing (if enabled) ---
  if (opts.isPostProcessEnable) {
    console.log('Applying CPU-based connected-component labeling...');
    const bwStartTime = performance.now();
    const BWInstance = new BWLabeler();

    // When true, components dropped by the per-class "largest component" filters
    // are repainted with the most common surviving label on their boundary
    // instead of becoming background. Costs one extra linear (6-neighbour) pass
    // over the volume; only applies to per-class filtering paths (not the binary
    // brain-mask path, where dropped voxels are genuinely outside the brain).
    // Default ON for the 104-class DK-atlas models (ids 5/14), where punching
    // background holes into a dense parcellation is rarely what we want; the
    // global opts flag can additionally enable it for other per-class models.
    const FILL_NEIGHBOR_DEFAULT_IDS = [5, 14];
    const relabelSuppressed =
      !!opts.fillSuppressedWithNeighborLabel || FILL_NEIGHBOR_DEFAULT_IDS.includes(modelEntry.id);

    // --- Noise guard -------------------------------------------------------
    // A garbage/noise segmentation fragments into a huge number of tiny
    // connected components (we've seen ~1.5M). The labeling pass itself stays
    // fast, but the per-component filtering below (O(cl^2) scans, per-class
    // Maps and sorts) would then freeze the tab for minutes. So we run ONE
    // cheap plain labeling pass first (no filtering), and if the component
    // count is implausibly high we abort before any heavy work. Healthy
    // segmentations yield hundreds-to-few-thousand components, far below the
    // cap; noise yields ~1M+, far above it.
    const totalVoxels = Vshape[0] * Vshape[1] * Vshape[2];
    const NOISE_COMPONENT_CAP = Math.max(100000, Math.floor(totalVoxels * 0.01));
    const [guardComponentCount, guardLabels] = BWInstance.bwlabel(segmentationData, Vshape, 6, false, false);
    if (FORCE_SEGMENTATION_NOISE || guardComponentCount > NOISE_COMPONENT_CAP) {
      const msg =
        `Segmentation produced noise: ${guardComponentCount.toLocaleString()} ` +
        `disconnected regions (cap ${NOISE_COMPONENT_CAP.toLocaleString()}). ` +
        `The model output is unusable, so post-processing was aborted. ` +
        `Try re-running, switching backend (WebGPU/WebGL2), or another model.`;
      console.error('[postprocess] ' + msg);
      const err = new Error(msg);
      err.code = 'SEGMENTATION_NOISE';
      throw err;
    }

    let binarize = false;
    let onlyLargest = false;
    // Binary brain extraction/masking always keeps one connected foreground
    // component, regardless of its menu id.
    if (modelEntry.type === 'Brain_Extraction' || modelEntry.type === 'Brain_Masking') {
      binarize = true;
      onlyLargest = true;
    } else if ([1, 7].includes(modelEntry.id)) {
      // 3-class (1, 7): keep top-2 components per class (preserves a detached
      // cerebellum) but with a small size floor (drops tiny stray blobs).
      // Handled manually below.
      binarize = false;
      onlyLargest = false; // We handle filtering manually
    } else if ([5, 14].includes(modelEntry.id)) {
      // 104-class (5 = old, 14 = new deep DK-atlas): Use Strict Largest.
      // binarize=false + onlyLargest=true keeps the largest connected component
      // PER CLASS (not a single binarized blob), which is what a multi-region
      // atlas parcellation needs.
      binarize = false;
      onlyLargest = true;
    } else if ([3, 8, 9].includes(modelEntry.id)) {
      // 18-class: Mixed logic (targets only)
      binarize = false;
      onlyLargest = false;
    } else {
      // 50-class (4), Brain Extractions, Masks, etc.
      // Default legacy behavior: Binarize first, then keep largest blob.
      binarize = true;
      onlyLargest = true;
    }

    if ([1, 7].includes(modelEntry.id)) {
      // 3-class rank logic: keep the top-2 connected components PER CLASS so a
      // legitimately detached structure (this model sometimes segments the
      // cerebellum separately) is preserved. The size floor then drops any
      // kept component below SMALL_COMPONENT_MIN_RATIO of its class's largest,
      // removing tiny stray blobs (e.g. a misclassified chin speck) without
      // touching the much-larger cerebellum. Tune UP if noise remains; tune
      // DOWN if a small detached cerebellum ever gets clipped.
      const SMALL_COMPONENT_MIN_RATIO = 0.02;
      // Spatial gate (SURFACE distance in VOXELS): the 2nd component per class is
      // kept ONLY if its shortest empty-space path to the main brain is within
      // this many voxels. A detached cerebellum nearly touches the cerebrum
      // (surfDist ~= a few voxels of CSF) and is kept; a phantom blob outside the
      // head is separated by a large empty gap and is dropped. We measure SURFACE
      // distance (multi-source BFS from the brain), NOT a bounding-box gap: the
      // cerebrum bbox fills most of a 256^3 head and encloses a far phantom's
      // bbox, giving a misleading gap of 0. Tune UP to keep farther pieces, DOWN
      // to be stricter. (BFS scans ~GAP+4 layers, so larger values cost a bit more.)
      const NEAR_BRAIN_MAX_GAP = 8;
      const DIAG_RANK_FILTER = false; // set true to log brain bbox + per-component surfDist/keep decision
      // Reuse the labeling already computed by the noise guard (identical call:
      // conn=6, binarize=false, onlyLargest=false) to avoid a redundant pass.
      const cl = guardComponentCount;
      const ls = guardLabels;
      const [_mx, filtered] = BWInstance.filter_clusters_by_rank(segmentationData, cl, ls, 2, SMALL_COMPONENT_MIN_RATIO, Vshape, relabelSuppressed, NEAR_BRAIN_MAX_GAP, DIAG_RANK_FILTER);
      segmentationData.set(filtered);
    } else if (!onlyLargest && [3, 8, 9].includes(modelEntry.id)) {
      // Mixed case (18-class) - Hierarchical approach:
      // Step 1: Binary largest connected component to establish brain boundary
      // This removes all disconnected noise/artifacts in one sweep
      const [_cl1, binaryMask] = BWInstance.bwlabel(segmentationData, Vshape, 6, true, true);

      // Apply binary mask to remove all disconnected voxels
      for (let i = 0; i < segmentationData.length; i++) {
        segmentationData[i] *= binaryMask[i];
      }

      // Step 2: Per-class filtering on the cleaned data
      // Get raw components from the masked (cleaned) segmentation
      const [cl, ls] = BWInstance.bwlabel(segmentationData, Vshape, 6, false, false);

      // Calculate targets (Explicit IDs for 18-class model)
      // 1: Cerebral-White-Matter
      // 2: Cerebral-Cortex
      // 5: Cerebellum-White-Matter
      // 6: Cerebellum-Cortex
      // 13: Brain-Stem
      const targetClasses = new Set([1, 2, 5, 6, 13]);

      // Apply mixed filter to keep only largest component for target classes
      const [_mx, filtered] = BWInstance.filter_clusters(segmentationData, cl, ls, targetClasses, Vshape, relabelSuppressed);
      segmentationData.set(filtered);

    } else if (!binarize && onlyLargest) {
      // Per-Class Strict (e.g. 104-class DK-atlas, ids 5/14): keep the largest
      // connected component per class. This is equivalent to
      // bwlabel(..., binarize=false, onlyLargest=true), but we reuse the
      // labeling already produced by the noise guard (same conn=6/binarize=false
      // pass) and only run the largest-cluster selection, avoiding a second
      // full relabeling of this perf-sensitive model.
      // DIAGNOSTIC (opt-in, output-neutral): inspect components on the ORIGINAL
      // segmentation (segmentationData not yet overwritten; guardLabels match it)
      // to understand why a region survives — see diagnose_components.
      if (opts.diagnoseEnclosedComponents) {
        BWInstance.diagnose_components(segmentationData, guardComponentCount, guardLabels, Vshape, { label: `model${modelEntry.id}`, topN: 60 });
      }
      const [_mx, res2] = BWInstance.largest_original_cluster_labels(segmentationData, guardComponentCount, guardLabels, Vshape, relabelSuppressed);
      segmentationData.set(res2);
    } else {
      // Standard cases (Legacy binarize-then-largest)
      const [_res1, res2] = BWInstance.bwlabel(segmentationData, Vshape, 6, binarize, onlyLargest);

      if (binarize) {
        // res2 is binary mask (0/1). Apply to segmentationData to clean it.
        // This preserves original classes inside the mask.
        for (let i = 0; i < segmentationData.length; i++) {
          segmentationData[i] *= res2[i];
        }
      } else {
        // res2 contains the filtered class IDs directly.
        // We replace segmentationData with it.
        segmentationData.set(res2);
      }
    }

    const bwTime = ((performance.now() - bwStartTime) / 1000).toFixed(4);
    console.log(`Connected-component labeling took: ${bwTime} seconds.`);
  }

  // --- Step 3: Apply Final Model Logic on the CPU ---
  switch (modelEntry.type) {
    case 'Brain_Masking': {
      const brainMask = new Uint8Array(segmentationData.length);
      for (let i = 0; i < segmentationData.length; i++) {
        brainMask[i] = segmentationData[i] !== 0 ? 1 : 0;
      }
      return brainMask;
    }
    case 'Brain_Extraction': {
      if (modelEntry.returnMaskForExtraction) {
        const brainMask = new Uint8Array(segmentationData.length);
        for (let i = 0; i < segmentationData.length; i++) {
          brainMask[i] = segmentationData[i] !== 0 ? 1 : 0;
        }
        return brainMask;
      }
      // Preserve the input NIfTI datatype. Rodent samples are float32; forcing
      // them through Uint8Array would collapse their 0..1 intensities.
      const maskedData = new niftiImage.constructor(segmentationData.length);
      for (let i = 0; i < segmentationData.length; i++) {
        const maskValue = segmentationData[i] !== 0 ? 1 : 0;
        maskedData[i] = niftiImage[i] * maskValue;
      }
      return maskedData;
    }
    default: {
      // For other cases, return the (potentially modified) segmentationData.
      // Make sure to return correct TypedArray (Uint8Array usually expected by display)
      // Original code returned Uint8Array. 
      // segmentationData is Int32Array (from tensor download).
      return new Uint8Array(segmentationData);
    }
  }
}

/**
 * Estimates the maximum number of elements in any intermediate tensor of the model.
 * @param {tf.LayersModel} model The TensorFlow.js model.
 * @param {number[]} inputShape The shape of the input tensor (including batch size).
 * @returns {number} The maximum number of elements.
 */
/**
 * Estimates the maximum number of elements in any intermediate tensor of the model.
 * Performs a precise calculation of peak memory usage by summing Input and Output sizes
 * for each layer invocation, assuming sequential execution.
 * 
 * @param {tf.LayersModel} model The TensorFlow.js model.
 * @param {number[]} inputShape The shape of the input tensor (including batch size).
 * @param {boolean} isChannelLast Whether the model uses channels-last data format.
 * @returns {number} The maximum number of elements.
 */
export function estimateMaxIntermediateTensorSize(model, inputShape, isChannelLast) {
  let maxElements = 0;

  // 1. Calculate Spatial Volume (assumed constant/upper-bound size 256^3)
  let spatialVol = 1;
  // Heuristic: multiply the middle dimensions.
  if (isChannelLast) {
    // Expected: [Batch, D, H, W, C]
    if (inputShape.length === 5) {
      spatialVol = inputShape[1] * inputShape[2] * inputShape[3];
    } else {
      // Fallback usually [1, D, H, W]
      for (let i = 0; i < inputShape.length; i++) {
        if (inputShape[i] > 1) spatialVol *= inputShape[i];
      }
    }
  } else {
    // Expected: [Batch, C, D, H, W]
    if (inputShape.length === 5) {
      spatialVol = inputShape[2] * inputShape[3] * inputShape[4];
    } else {
      for (let i = 0; i < inputShape.length; i++) {
        if (inputShape[i] > 32) spatialVol *= inputShape[i];
      }
    }
  }

  // 2. Iterate Layers to find Peak Memory (Input + Output) and Max Output
  // checkMemoryAllocation tests both packed (peak) and unpacked (maxOutput).
  let maxOutputElements = 0;
  // Largest SINGLE tensor, which is the right quantity for a TEXTURE-DIMENSION
  // check. `maxElements` below is input+output summed, which is a fair proxy for
  // total GPU MEMORY but is wrong for the texture limit: the input activation and
  // the output activation are two separate tensors in two separate textures, and
  // MAX_TEXTURE_SIZE applies per texture. Summing them inflates the computed
  // dimension by sqrt(2), which is enough to push a 16-channel layer on a
  // 256x204x204 crop from 6528 to 9232 -- under Chrome's 16384 either way, but
  // straddling Firefox's 8192. That is why this only ever bit on Firefox.
  let maxSingleElements = 0;

  if (model && model.layers) {
    const _numLayers = model.layers.length;
    for (let _li = 0; _li < _numLayers; _li++) {
      const layer = model.layers[_li];
      const _isFinalLayer = (_li === _numLayers - 1);
      // A. Calculate Output Channels
      let outputChannels = 0;
      let outputShape = layer.outputShape;
      // Normalize outputShape to array
      if (Array.isArray(outputShape) && Array.isArray(outputShape[0])) {
        outputShape = outputShape[0];
      }
      if (Array.isArray(outputShape)) {
        if (isChannelLast) {
          outputChannels = outputShape[outputShape.length - 1];
        } else {
          // [Null, C, D, H, W]
          outputChannels = outputShape[1];
        }
      }

      // B. Calculate Input Channels (Robust)
      let inputChannels = 0;

      // B1. Try batchInputShape
      const inputShapes = layer.batchInputShape;
      const getChannelsFromShape = (s) => {
        if (!Array.isArray(s)) return 0;
        if (isChannelLast) return s[s.length - 1];
        return s[1];
      };

      if (inputShapes) {
        if (Array.isArray(inputShapes) && Array.isArray(inputShapes[0])) {
          // Array of shapes (e.g. Concatenate layer inputs)
          for (const s of inputShapes) {
            inputChannels += getChannelsFromShape(s);
          }
        } else if (Array.isArray(inputShapes)) {
          // Single shape
          inputChannels = getChannelsFromShape(inputShapes);
        }
      }

      // B2. Fallback: Try Weights (Kernels)
      // Conv3D kernel: [D, H, W, In, Out]
      if (inputChannels === 0 && layer.weights && layer.weights.length > 0) {
        // We assume the first weight is the kernel.
        // Warning: Accessing .val or .tensor might be expensive? 
        // layer.weights is array of LayerVariable. variable.shape is available.
        const w = layer.weights[0];
        if (w && w.shape) {
          if (w.shape.length === 5) { // Conv3D
            // Kernel layout is channel selection logic dependent?
            // TF.js Conv3D kernel is [D, H, W, In, Out] usually.
            inputChannels = w.shape[3];
          } else if (w.shape.length === 4) { // Conv2D
            // [H, W, In, Out]
            inputChannels = w.shape[2];
          }
        }
      }

      // B3. Final Fallback: Assume Input = Output (e.g. Activation layers)
      if (inputChannels === 0) {
        inputChannels = outputChannels;
      }

      // C. Calculate Peak and Update Max
      if (typeof outputChannels === 'number' && typeof inputChannels === 'number') {
        const currentPeakElements = spatialVol * (inputChannels + outputChannels);
        const currentOutputElements = spatialVol * outputChannels;

        // The FINAL layer (e.g. 24->104) is never run densely in the fast loop:
        // it is always handled separately by chunkedArgMax / SequentialConvLayer.
        // So its large channel count must NOT inflate the "peak intermediate"
        // used to decide full SeqConv -- otherwise any head cropping above
        // ~205^3 wrongly trips full SeqConv (24 sliced conv3d + concats per
        // layer) and the model crawls. The unpacked maxOutput check below still
        // accounts for the final layer and routes it to chunkedArgMax.
        if (!_isFinalLayer && currentPeakElements > maxElements) {
          maxElements = currentPeakElements;
        }
        // Same exclusion of the final layer, for the same reason: it is always
        // routed to chunkedArgMax / SequentialConvLayer and never materialized
        // densely, so its channel count must not decide the backbone's path.
        const currentSingleElements = spatialVol * Math.max(inputChannels, outputChannels);
        if (!_isFinalLayer && currentSingleElements > maxSingleElements) {
          maxSingleElements = currentSingleElements;
        }
        // Track output of LAST layer (overwrite each iteration)
        // Only the FINAL output is unpacked for argmax. Intermediates are packed.
        maxOutputElements = currentOutputElements;
      }
    }
  }

  // Fallback
  if (maxElements === 0) {
    maxElements = spatialVol * 32 * 2;
  }
  if (maxSingleElements === 0) {
    maxSingleElements = spatialVol * 32;
  }

  // Does the largest intermediate get UNPACKED at some point? For the affine-GN
  // and GN models it does, and that changes the limit by 2x in the dimension.
  // LayerNormInPlace does `x.transpose([0,4,1,2,3]).reshape([C, N])`, which
  // materializes the whole activation as a flat [C, N] tensor with no packing --
  // 1 element per texel, so the texture is ceil(sqrt(N)) on a side, not
  // ceil(sqrt(N/4)). tfjs squarifies it, which is where the "[12481x12481]" in
  // the failure message comes from: 15 * 204*202*252 = 155766240 elements.
  //
  // This is measured from the failure, not assumed: 12481 = ceil(sqrt(15V))
  // exactly, and 6241 = ceil(sqrt(15V/4)) is what the packed check computed.
  const hasUnpackedIntermediate = !!(model && model.layers &&
    model.layers.some((l) => typeof l.name === 'string' && l.name.endsWith('_gn')));

  // Return BOTH peak (In+Out) and max output separately
  // This allows checkMemoryAllocation to test packed (peak) and unpacked (maxOutput)
  console.log(`[Estimator] Total Layers: ${model?.layers?.length}, Peak(in+out): ${maxElements}, ` +
    `MaxSingle: ${maxSingleElements}, Final Output: ${maxOutputElements}, ` +
    `unpackedIntermediate: ${hasUnpackedIntermediate}`);
  return {
    peak: maxElements,
    maxSingle: maxSingleElements,
    maxOutput: maxOutputElements,
    hasUnpackedIntermediate,
  };
}

/**
 * Proactively checks if a tensor of the specified size can be allocated on the GPU.
 * @param {number} peakElements Total elements for peak allocation (In+Out) - used for packed check.
 * @param {number} maxOutputElements Elements for largest single output - used for unpacked check.
 * @returns {boolean} True if allocation is likely to succeed, false otherwise.
 */
export function checkMemoryAllocation(peakElements, maxOutputElements) {
  try {
    const backend = tf.backend();
    if (backend && backend.gpgpu && backend.gpgpu.gl) {
      const maxTextureSize = backend.gpgpu.gl.getParameter(backend.gpgpu.gl.MAX_TEXTURE_SIZE);

      // TF.js WebGL uses PACKED textures (4 elem/pixel) for INTERMEDIATE layers.
      // However, the FINAL OUTPUT (e.g., class logits before argmax) is UNPACKED (1 elem/pixel).

      // Check for PACKED intermediates (4 elements per pixel)
      const packedNeededPixels = Math.ceil(peakElements / 4);
      const packedNeededDim = Math.ceil(Math.sqrt(packedNeededPixels));

      // Check for UNPACKED final output (1 element per pixel)
      const unpackedNeededPixels = maxOutputElements; // 1 elem/pixel
      const unpackedNeededDim = Math.ceil(Math.sqrt(unpackedNeededPixels));

      console.log(`[Memory Check] Peak: ${peakElements}, MaxOutput: ${maxOutputElements}, Packed Dim: ${packedNeededDim}, Unpacked Dim: ${unpackedNeededDim}, MaxTextureSize: ${maxTextureSize}`);

      if (packedNeededDim > maxTextureSize) {
        console.warn(`Proactive check (PACKED): Tensor size ${peakElements} requires approx ${packedNeededDim}x${packedNeededDim} texture. Exceeds MAX_TEXTURE_SIZE ${maxTextureSize}`);
        return false;
      }

      if (unpackedNeededDim > maxTextureSize) {
        console.warn(`Proactive check (UNPACKED): Max output ${maxOutputElements} requires approx ${unpackedNeededDim}x${unpackedNeededDim} texture. Exceeds MAX_TEXTURE_SIZE ${maxTextureSize}`);
        return false;
      }
    }
  } catch (e) {
    console.warn("Could not check texture size limits:", e);
  }

  // If we reach here, texture size check passed (or no backend available)
  return true;
}