{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "06270b57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import pathlib\n",
    "import numpy as np\n",
    "from scipy import ndimage\n",
    "import skimage.measure\n",
    "import dask.array as da\n",
    "import napari"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c2bbccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names_colors = {\n",
    "    \"bronchioles\": [150, 99, 23, 255],  # brown\n",
    "    \"alveoli\": [23, 80, 150, 255],  # dark blue\n",
    "    #\"alveoli\": [23, 80, 150, 0],  # dark blue\n",
    "    \"vasculature\": [150, 31, 23, 255],  # dark red\n",
    "    \"cancer\": [199, 196, 147, 255],  # v dark purple\n",
    "    \"nonexpanded\": [23, 80, 150, 255],  # dark blue\n",
    "    #\"nonexpanded\": [23, 80, 150, 0],  # dark blue\n",
    "    #\"whitespace\": [255, 255, 255, 255],  # white\n",
    "    \"whitespace\": [255, 255, 255, 0],  # white\n",
    "    \"collagen\": [242, 167, 227, 255],  # light pink\n",
    "    #\"collagen\": [242, 167, 227, 0],  # light pink\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4b3114b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "registered_labels_zarr_path = (\n",
    "    pathlib.Path().resolve().parent.parent\n",
    "    / \"testing_python_FULL_20250929\"\n",
    "    / \"registered_class_labels.zarr\"\n",
    "    # / \"registered_class_labels_2_mpp.zarr\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "695184c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/margareteminizer/miniconda3/envs/coda-scratch/lib/python3.10/site-packages/zarr/creation.py:614: UserWarning: ignoring keyword argument 'read_only'\n",
      "  compressor, fill_value = _kwargs_compat(compressor, fill_value, kwargs)\n"
     ]
    }
   ],
   "source": [
    "labels_data = da.from_zarr(str(registered_labels_zarr_path/\"arr_0\"))\n",
    "labels_data = labels_data.transpose(2, 0, 1)  # Change shape from (y, x, z) to (z, y, x)\n",
    "labels_data = labels_data.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "330633ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-10-10 13:25:25] Processing class bronchioles (1)\n",
      "[2025-10-10 13:25:25]\tLabeling connected components...\n",
      "[2025-10-10 13:25:27]\tComputing component sizes...\n",
      "[2025-10-10 13:25:30]\tKeeping 3 out of 60413 components\n",
      "[2025-10-10 13:25:30]\tFiltering labels...\n",
      "[2025-10-10 13:25:32] Processing class alveoli (2)\n",
      "[2025-10-10 13:25:32]\tLabeling connected components...\n",
      "[2025-10-10 13:25:36]\tComputing component sizes...\n",
      "[2025-10-10 13:25:38]\tKeeping 1 out of 973122 components\n",
      "[2025-10-10 13:25:38]\tFiltering labels...\n",
      "[2025-10-10 13:25:39] Processing class vasculature (3)\n",
      "[2025-10-10 13:25:39]\tLabeling connected components...\n",
      "[2025-10-10 13:25:42]\tComputing component sizes...\n",
      "[2025-10-10 13:25:45]\tKeeping 511 out of 269294 components\n",
      "[2025-10-10 13:25:45]\tFiltering labels...\n",
      "[2025-10-10 13:25:45] Processing class cancer (4)\n",
      "[2025-10-10 13:25:45]\tLabeling connected components...\n",
      "[2025-10-10 13:25:47]\tComputing component sizes...\n",
      "[2025-10-10 13:25:50]\tKeeping 8753 out of 16597 components\n",
      "[2025-10-10 13:25:50]\tFiltering labels...\n",
      "[2025-10-10 13:25:50] Processing class nonexpanded (5)\n",
      "[2025-10-10 13:25:50]\tLabeling connected components...\n",
      "[2025-10-10 13:25:53]\tComputing component sizes...\n",
      "[2025-10-10 13:25:55]\tKeeping 5 out of 7277 components\n",
      "[2025-10-10 13:25:55]\tFiltering labels...\n",
      "[2025-10-10 13:25:56] Processing class collagen (7)\n",
      "[2025-10-10 13:25:56]\tLabeling connected components...\n",
      "[2025-10-10 13:25:58]\tComputing component sizes...\n",
      "[2025-10-10 13:26:01]\tKeeping 19 out of 202782 components\n",
      "[2025-10-10 13:26:01]\tFiltering labels...\n"
     ]
    }
   ],
   "source": [
    "import datetime\n",
    "from numba import njit, prange\n",
    "\n",
    "MIN_FRAC = 0.0005  # Minimum fraction of the largest component to keep\n",
    "\n",
    "\n",
    "def get_tss() -> str:\n",
    "    return datetime.datetime.now().strftime(\"[%Y-%m-%d %H:%M:%S]\")\n",
    "\n",
    "\n",
    "@njit(parallel=True)\n",
    "def _filter_labels_numba(labeled, keep_labels):\n",
    "    keep = np.zeros(labeled.shape, dtype=np.bool_)\n",
    "    keep_set = set(keep_labels)\n",
    "    flat = labeled.ravel()\n",
    "    keep_flat = keep.ravel()\n",
    "    for i in prange(flat.size):\n",
    "        if flat[i] in keep_set and flat[i] != 0:\n",
    "            keep_flat[i] = True\n",
    "    return keep\n",
    "\n",
    "\n",
    "filtered_labels = np.zeros_like(labels_data, dtype=labels_data.dtype)\n",
    "for class_idx, class_name in enumerate(class_names_colors.keys(), start=1):\n",
    "    if class_name in (\"whitespace\",):\n",
    "        continue\n",
    "    print(f\"{get_tss()} Processing class {class_name} ({class_idx})\")\n",
    "    mask = labels_data == class_idx\n",
    "    if not np.any(mask):\n",
    "        continue\n",
    "    # Label connected components\n",
    "    print(f\"{get_tss()}\\tLabeling connected components...\")\n",
    "    labeled, n_labels = ndimage.label(mask)\n",
    "    if n_labels == 0:\n",
    "        continue\n",
    "    # Compute sizes of each component\n",
    "    print(f\"{get_tss()}\\tComputing component sizes...\")\n",
    "    sizes = np.bincount(labeled.ravel())\n",
    "    sizes[0] = 0  # background\n",
    "    max_size = sizes.max()\n",
    "    if max_size == 0:\n",
    "        continue\n",
    "    # Find which labels to keep\n",
    "    keep_labels = np.where(sizes >= MIN_FRAC * max_size)[0]\n",
    "    print(f\"{get_tss()}\\tKeeping {len(keep_labels)} out of {n_labels} components\")\n",
    "    # Use numba to filter in parallel\n",
    "    print(f\"{get_tss()}\\tFiltering labels...\")\n",
    "    filtered_class = _filter_labels_numba(labeled, keep_labels)\n",
    "    filtered_labels[filtered_class] = class_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5c4067c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/margareteminizer/miniconda3/envs/coda-scratch/lib/python3.10/site-packages/napari/_vispy/layers/scalar_field.py:198: UserWarning: data shape (150, 4861, 4071) exceeds GL_MAX_TEXTURE_SIZE 2048 in at least one axis and will be downsampled. Rendering is currently in 3D mode.\n",
      "  warnings.warn(\n",
      "/Users/margareteminizer/miniconda3/envs/coda-scratch/lib/python3.10/site-packages/napari/_vispy/layers/scalar_field.py:198: UserWarning: data shape (150, 4861, 4071) exceeds GL_MAX_TEXTURE_SIZE 2048 in at least one axis and will be downsampled. Rendering is currently in 3D mode.\n",
      "  warnings.warn(\n",
      "/Users/margareteminizer/miniconda3/envs/coda-scratch/lib/python3.10/site-packages/napari/_vispy/layers/scalar_field.py:198: UserWarning: data shape (150, 4861, 4071) exceeds GL_MAX_TEXTURE_SIZE 2048 in at least one axis and will be downsampled. Rendering is currently in 3D mode.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "colors = {\n",
    "    i: tuple([v / 255.0 for v in color])\n",
    "    for i, color in enumerate(class_names_colors.values(), start=1)\n",
    "}\n",
    "colors[0] = (0, 0, 0, 0)  # Make background transparent\n",
    "colors[None] = (0, 0, 0, 0)  # Also make None transparent\n",
    "\n",
    "labels = dict(enumerate(class_names_colors.keys(), start=1))\n",
    "labels[0] = \"unassigned\"\n",
    "\n",
    "viewer = napari.Viewer()\n",
    "viewer.add_labels(\n",
    "    labels_data,\n",
    "    name=\"Registered Class Labels\",\n",
    "    colormap=colors,\n",
    "    axis_labels=(\"z\", \"y\", \"x\"),\n",
    "    opacity=1.0,\n",
    "    scale=(2.0, 1.0, 1.0),\n",
    ")\n",
    "viewer.add_labels(\n",
    "    filtered_labels,\n",
    "    name=\"Filtered Class Labels\",\n",
    "    colormap=colors,\n",
    "    axis_labels=(\"z\", \"y\", \"x\"),\n",
    "    opacity=1.0,\n",
    "    scale=(2.0, 1.0, 1.0),\n",
    ")\n",
    "napari.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5dd373ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface_meshes = {}\n",
    "for i_class, class_name in enumerate(class_names_colors.keys(), start=1):\n",
    "    if class_name in (\"whitespace\", \"alveoli\", \"nonexpanded\"):\n",
    "        continue  # Skip whitespace class\n",
    "    # class_mask = labels_data == i_class\n",
    "    class_mask = filtered_labels == i_class\n",
    "    verts, faces, _, _ = skimage.measure.marching_cubes(class_mask, level=0.5)\n",
    "    surface_meshes[class_name] = (verts, faces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "150ce1c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = napari.Viewer()\n",
    "for class_name, (verts, faces) in surface_meshes.items():\n",
    "    color = tuple([v / 255.0 for v in class_names_colors[class_name]])\n",
    "    viewer.add_surface((verts, faces), name=class_name, colormap=color)\n",
    "napari.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd848e2b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "coda-scratch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}