main.js

import { Niivue } from "@niivue/niivue";
import { runInference } from "./brainchop-mainthread.js";
import { inferenceModelsList, brainChopOpts } from "./brainchop-parameters.js";
import { isChrome, localSystemDetails } from "./brainchop-diagnostics.js";
import MyWorker from "./brainchop-webworker.js?worker";

async function main() {
  dragMode.onchange = async function () {
    nv1.opts.dragMode = this.selectedIndex;
  };
  drawDrop.onchange = async function () {
    if (nv1.volumes.length < 2) {
      window.alert("No segmentation open (use the Segmentation pull down)");
      drawDrop.selectedIndex = -1;
      return;
    }
    if (!nv1.drawBitmap) {
      window.alert("No drawing (hint: use the Draw pull down to select a pen)");
      drawDrop.selectedIndex = -1;
      return;
    }
    const mode = parseInt(this.value);
    if (mode === 0) {
      nv1.drawUndo();
      drawDrop.selectedIndex = -1;
      return;
    }
    let img = nv1.volumes[1].img;
    let draw = await nv1.saveImage({ filename: "", isSaveDrawing: true });
    const niiHdrBytes = 352;
    const nvox = draw.length;
    if (mode === 1) {
      //append
      for (let i = 0; i < nvox; i++) if (draw[niiHdrBytes + i] > 0) img[i] = 1;
    }
    if (mode === 2) {
      //delete
      for (let i = 0; i < nvox; i++) if (draw[niiHdrBytes + i] > 0) img[i] = 0;
    }
    nv1.closeDrawing();
    nv1.updateGLVolume();
    nv1.setDrawingEnabled(false);
    penDrop.selectedIndex = -1;
    drawDrop.selectedIndex = -1;
  };
  penDrop.onchange = async function () {
    const mode = parseInt(this.value);
    nv1.setDrawingEnabled(mode >= 0);
    if (mode >= 0) nv1.setPenValue(mode & 7, mode > 7);
  };
  aboutBtn.onclick = function () {
    window.alert(
      "Drag and drop NIfTI images. Use pulldown menu to choose brainchop model",
    );
  };
  diagnosticsBtn.onclick = function () {
    if (diagnosticsString.length < 1) {
      window.alert(
        "No diagnostic string generated: run a model to create diagnostics",
      );
      return;
    }
    missingLabelStatus = missingLabelStatus.slice(0, -2);
    if (missingLabelStatus !== "") {
      if (diagnosticsString.includes('Status: OK')) {
        diagnosticsString = diagnosticsString.replace('Status: OK', `Status: ${missingLabelStatus}`);
      }
    }
    missingLabelStatus = ""
    navigator.clipboard.writeText(diagnosticsString);
    window.alert("Diagnostics copied to clipboard\n" + diagnosticsString);
  };
  opacitySlider0.oninput = function () {
    nv1.setOpacity(0, opacitySlider0.value / 255);
    nv1.updateGLVolume();
  };
  opacitySlider1.oninput = function () {
    nv1.setOpacity(1, opacitySlider1.value / 255);
  };
  async function ensureConformed() {
    const nii = nv1.volumes[0];
    let isConformed =
      nii.dims[1] === 256 && nii.dims[2] === 256 && nii.dims[3] === 256;
    if (
      nii.permRAS[0] !== -1 ||
      nii.permRAS[1] !== 3 ||
      nii.permRAS[2] !== -2
    ) {
      isConformed = false;
    }
    if (isConformed) {
      return;
    }
    const nii2 = await nv1.conform(nii, false);
    await nv1.removeVolume(nv1.volumes[0]);
    await nv1.addVolume(nii2);
  }
  async function closeAllOverlays() {
    while (nv1.volumes.length > 1) {
      await nv1.removeVolume(nv1.volumes[1]);
    }
  }
  modelSelect.onchange = async function () {
    if (this.selectedIndex < 0) {
      modelSelect.selectedIndex = 11;
    }
    await closeAllOverlays();
    await ensureConformed();
    const model = inferenceModelsList[this.selectedIndex];
    const opts = brainChopOpts;
    // opts.rootURL should be the url without the query string
    const urlParams = new URL(window.location.href);
    // remove the query string
    opts.rootURL = urlParams.origin + urlParams.pathname;
    const isLocalhost = Boolean(
      window.location.hostname === "localhost" ||
        // [::1] is the IPv6 localhost address.
        window.location.hostname === "[::1]" ||
        // 127.0.0.1/8 is considered localhost for IPv4.
        window.location.hostname.match(
          /^127(?:\.(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)){3}$/,
        ),
    );
    if (isLocalhost) {
      opts.rootURL = location.protocol + "//" + location.host;
    }
    if (workerCheck.checked) {
      if (typeof chopWorker !== "undefined") {
        console.log(
          "Unable to start new segmentation: previous call has not completed",
        );
        return;
      }
      chopWorker = await new MyWorker({ type: "module" });
      const hdr = {
        datatypeCode: nv1.volumes[0].hdr.datatypeCode,
        dims: nv1.volumes[0].hdr.dims,
      };
      const msg = {
        opts,
        modelEntry: model,
        niftiHeader: hdr,
        niftiImage: nv1.volumes[0].img,
      };
      chopWorker.postMessage(msg);
      chopWorker.onmessage = function (event) {
        const cmd = event.data.cmd;
        if (cmd === "ui") {
          if (event.data.modalMessage !== "") {
            chopWorker.terminate();
            chopWorker = undefined;
          }
          callbackUI(
            event.data.message,
            event.data.progressFrac,
            event.data.modalMessage,
            event.data.statData,
          );
        }
        if (cmd === "img") {
          chopWorker.terminate();
          chopWorker = undefined;
          callbackImg(event.data.img, event.data.opts, event.data.modelEntry);
        }
      };
    } else {
      runInference(
        opts,
        model,
        nv1.volumes[0].hdr,
        nv1.volumes[0].img,
        callbackImg,
        callbackUI,
      );
    }
  };
  saveImgBtn.onclick = function () {
    nv1.volumes[1].saveToDisk("segmentaion.nii.gz");
  };
  saveSceneBtn.onclick = function () {
    nv1.saveDocument("brainchop.nvd");
  };
  workerCheck.onchange = function () {
    modelSelect.onchange();
  };
  clipCheck.onchange = function () {
    if (clipCheck.checked) {
      nv1.setClipPlane([0, 0, 90]);
    } else {
      nv1.setClipPlane([2, 0, 90]);
    }
  };
  function doLoadImage() {
    opacitySlider0.oninput();
  }
  async function fetchJSON(fnm) {
    const response = await fetch(fnm);
    const js = await response.json();
    return js;
  }
  async function getUniqueValuesAndCounts(uint8Array) {
    // Use a Map to count occurrences
    const countsMap = new Map();

    for (let i = 0; i < uint8Array.length; i++) {
      const value = uint8Array[i];
      if (countsMap.has(value)) {
        countsMap.set(value, countsMap.get(value) + 1);
      } else {
        countsMap.set(value, 1);
      }
    }

    // Convert the Map to an array of objects
    const result = Array.from(countsMap, ([value, count]) => ({
      value,
      count,
    }));

    return result;
  }
  async function createLabeledCounts(uniqueValuesAndCounts, labelStrings) {
    if (uniqueValuesAndCounts.length !== labelStrings.length) {
      missingLabelStatus = "Failed to Predict Labels - "
      console.error(
        "Mismatch in lengths: uniqueValuesAndCounts has",
        uniqueValuesAndCounts.length,
        "items, but labelStrings has",
        labelStrings.length,
        "items.",
      );
    }

    return labelStrings.map((label, index) => {
      // Find the entry matching the current label index
      const entry = uniqueValuesAndCounts.find(item => item.value === index);

      // If an entry is found, append the count value with 'mm3', otherwise show 'Missing'
      const countText = entry ? `${entry.count} mm3` : "Missing";

      countText === "Missing"
      ? missingLabelStatus += `${label}, ` : null;

      return `${label}   ${countText}`;
    });
  }
  async function callbackImg(img, opts, modelEntry) {
    closeAllOverlays();
    const overlayVolume = await nv1.volumes[0].clone();
    overlayVolume.zeroImage();
    overlayVolume.hdr.scl_inter = 0;
    overlayVolume.hdr.scl_slope = 1;
    overlayVolume.img = new Uint8Array(img);
    const roiVolumes = await getUniqueValuesAndCounts(overlayVolume.img);
    console.log(roiVolumes);
    if (modelEntry.colormapPath) {
      const cmap = await fetchJSON(modelEntry.colormapPath);
      const newLabels = await createLabeledCounts(roiVolumes, cmap["labels"]);
      console.log(newLabels);
      overlayVolume.setColormapLabel({
        R: cmap["R"],
        G: cmap["G"],
        B: cmap["B"],
        labels: newLabels,
      });
      // n.b. most models create indexed labels, but those without colormap mask scalar input
      overlayVolume.hdr.intent_code = 1002; // NIFTI_INTENT_LABEL
    } else {
      let colormap = opts.atlasSelectedColorTable.toLowerCase();
      const cmaps = nv1.colormaps();
      if (!cmaps.includes(colormap)) {
        colormap = "actc";
      }
      overlayVolume.colormap = colormap;
    }
    overlayVolume.opacity = opacitySlider1.value / 255;
    await nv1.addVolume(overlayVolume);
  }
  async function reportTelemetry(statData) {
    if (typeof statData === "string" || statData instanceof String) {
      function strToArray(str) {
        const list = JSON.parse(str);
        const array = [];
        for (const key in list) {
          array[key] = list[key];
        }
        return array;
      }
      statData = strToArray(statData);
    }
    statData = await localSystemDetails(statData, nv1.gl);
    diagnosticsString =
      ":: Diagnostics can help resolve issues https://github.com/neuroneural/brainchop/issues ::\n";
    for (const key in statData) {
      diagnosticsString += key + ": " + statData[key] + "\n";
    }
  }
  function callbackUI(
    message = "",
    progressFrac = -1,
    modalMessage = "",
    statData = [],
  ) {
    if (message !== "") {
      console.log(message);
      document.getElementById("location").innerHTML = message;
    }
    if (isNaN(progressFrac)) {
      // memory issue
      memstatus.style.color = "red";
      memstatus.innerHTML = "Memory Issue";
    } else if (progressFrac >= 0) {
      modelProgress.value = progressFrac * modelProgress.max;
    }
    if (modalMessage !== "") {
      window.alert(modalMessage);
    }
    if (Object.keys(statData).length > 0) {
      reportTelemetry(statData);
    }
  }
  function handleLocationChange(data) {
    document.getElementById("location").innerHTML = data.string
      .split("   ")
      .map((value) => `<p style="font-size: 14px;margin:0px;">${value}</p>`)
      .join("");
  }
  const defaults = {
    backColor: [0.4, 0.4, 0.4, 1],
    show3Dcrosshair: true,
    onLocationChange: handleLocationChange,
  };
  let diagnosticsString = "";
  let missingLabelStatus = ""
  let chopWorker;
  const nv1 = new Niivue(defaults);
  nv1.attachToCanvas(gl1);
  nv1.opts.dragMode = nv1.dragModes.pan;
  nv1.opts.multiplanarForceRender = true;
  nv1.opts.yoke3Dto2DZoom = true;
  nv1.opts.crosshairGap = 11;
  nv1.setInterpolation(true);
  await nv1.loadVolumes([{ url: "./t1_crop.nii.gz" }]);
  for (let i = 0; i < inferenceModelsList.length; i++) {
    const option = document.createElement("option");
    option.text = inferenceModelsList[i].modelName;
    option.value = inferenceModelsList[i].id.toString();
    modelSelect.appendChild(option);
  }
  nv1.onImageLoaded = doLoadImage;
  modelSelect.selectedIndex = -1;
  drawDrop.selectedIndex = -1;
  workerCheck.checked = await isChrome(); // TODO: Safari does not yet support WebGL TFJS webworkers, test FireFox
  // uncomment next two lines to automatically run segmentation when web page is loaded
  // modelSelect.selectedIndex = 11
  // modelSelect.onchange()

  // get the query string parameter model.
  // if set, select the model from the dropdown list and call the modelSelect.onchange() function
  const urlParams = new URLSearchParams(window.location.search);
  const modelParam = urlParams.get("model");
  if (modelParam) {
    // make sure the model index is a number
    modelSelect.selectedIndex = Number(modelParam);
    modelSelect.onchange();
  }
}

async function updateStarCount() {
  try {
    const response = await fetch(
      `https://api.github.com/repos/neuroneural/brainchop`,
    );
    const data = await response.json();
    document.getElementById("star-count").textContent = data.stargazers_count;
  } catch (error) {
    console.error("Error fetching star count:", error);
  }
}
(async function () {
  await main();
  updateStarCount();
})();