run_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation and Requirements\n",
    "\n",
    "Please refer to the [_featureforest repo_](https://github.com/juglab/featureforest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from pathlib import Path\n",
    "\n",
    "import h5py\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image, ImageSequence\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "from tqdm.notebook import trange, tqdm\n",
    "\n",
    "from featureforest.models import get_available_models, get_model\n",
    "from featureforest.models.SAM import SAMAdapter\n",
    "from featureforest.utils.data import (\n",
    "    patchify,\n",
    "    is_image_rgb, get_stride_margin,\n",
    "    get_num_patches, get_stride_margin\n",
    ")\n",
    "from featureforest.postprocess import (\n",
    "    postprocess,\n",
    "    postprocess_with_sam, postprocess_with_sam_auto,\n",
    "    get_sam_auto_masks\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_slice_features(\n",
    "    image: np.ndarray,\n",
    "    patch_size: int,\n",
    "    overlap: int,\n",
    "    model_adapter,\n",
    "    storage_group,\n",
    "):\n",
    "    \"\"\"Extract the model features for one slice and save them into storage file.\"\"\"\n",
    "    # image to torch tensor\n",
    "    img_data = torch.from_numpy(image).to(torch.float32) / 255.0\n",
    "    # for sam the input image should be 4D: BxCxHxW ; an RGB image.\n",
    "    if is_image_rgb(image):\n",
    "        # it's already RGB, put the channels first and add a batch dim.\n",
    "        img_data = img_data[..., :3]  # ignore the Alpha channel (in case of PNG).\n",
    "        img_data = img_data.permute([2, 0, 1]).unsqueeze(0)\n",
    "    else:\n",
    "        img_data = img_data.unsqueeze(0).unsqueeze(0).expand(-1, 3, -1, -1)\n",
    "\n",
    "    # get input patches\n",
    "    data_patches = patchify(img_data, patch_size, overlap)\n",
    "    num_patches = len(data_patches)\n",
    "\n",
    "    # set a low batch size\n",
    "    batch_size = 8\n",
    "    # for big SAM we need even lower batch size :(\n",
    "    if isinstance(model_adapter, SAMAdapter):\n",
    "        batch_size = 2\n",
    "\n",
    "    num_batches = int(np.ceil(num_patches / batch_size))\n",
    "    # prepare storage for the slice embeddings\n",
    "    total_channels = model_adapter.get_total_output_channels()\n",
    "    stride, _ = get_stride_margin(patch_size, overlap)\n",
    "\n",
    "    if model_adapter.name not in storage_group:\n",
    "        dataset = storage_group.create_dataset(\n",
    "            model_adapter.name, shape=(num_patches, stride, stride, total_channels)\n",
    "        )\n",
    "    else:\n",
    "        dataset = storage_group[model_adapter.name]\n",
    "\n",
    "    # get sam encoder output for image patches\n",
    "    # print(\"\\nextracting slice features:\")\n",
    "    for b_idx in tqdm(range(num_batches), desc=\"extracting slice feature:\"):\n",
    "        # print(f\"batch #{b_idx + 1} of {num_batches}\")\n",
    "        start = b_idx * batch_size\n",
    "        end = start + batch_size\n",
    "        slice_features = model_adapter.get_features_patches(\n",
    "            data_patches[start:end].to(model_adapter.device)\n",
    "        )\n",
    "        if not isinstance(slice_features, tuple):\n",
    "            # model has only one output\n",
    "            num_out = slice_features.shape[0]  # to take care of the last batch size\n",
    "            dataset[start : start + num_out] = slice_features\n",
    "        else:\n",
    "            # model has more than one output: put them into storage one by one\n",
    "            ch_start = 0\n",
    "            for feat in slice_features:\n",
    "                num_out = feat.shape[0]\n",
    "                ch_end = ch_start + feat.shape[-1]  # number of features\n",
    "                dataset[start : start + num_out, :, :, ch_start:ch_end] = feat\n",
    "                ch_start = ch_end\n",
    "\n",
    "\n",
    "def predict_slice(\n",
    "    rf_model, patch_dataset, model_adapter,\n",
    "    img_height, img_width, patch_size, overlap\n",
    "):\n",
    "    \"\"\"Predict a slice patch by patch\"\"\"\n",
    "    segmentation_image = []\n",
    "    # shape: N x target_size x target_size x C\n",
    "    feature_patches = patch_dataset[:]\n",
    "    num_patches = feature_patches.shape[0]\n",
    "    total_channels = model_adapter.get_total_output_channels()\n",
    "    stride, margin = get_stride_margin(patch_size, overlap)\n",
    "\n",
    "    for i in tqdm(\n",
    "        range(num_patches), desc=\"Predicting slice patches\", position=1, leave=True\n",
    "    ):\n",
    "        input_data = feature_patches[i].reshape(-1, total_channels)\n",
    "        predictions = rf_model.predict(input_data).astype(np.uint8)\n",
    "        segmentation_image.append(predictions)\n",
    "\n",
    "    segmentation_image = np.vstack(segmentation_image)\n",
    "    # reshape into the image size + padding\n",
    "    patch_rows, patch_cols = get_num_patches(\n",
    "        img_height, img_width, patch_size, overlap\n",
    "    )\n",
    "    segmentation_image = segmentation_image.reshape(\n",
    "        patch_rows, patch_cols, stride, stride\n",
    "    )\n",
    "    segmentation_image = np.moveaxis(segmentation_image, 1, 2).reshape(\n",
    "        patch_rows * stride,\n",
    "        patch_cols * stride\n",
    "    )\n",
    "    # skip paddings\n",
    "    segmentation_image = segmentation_image[:img_height, :img_width]\n",
    "\n",
    "    return segmentation_image\n",
    "\n",
    "\n",
    "def apply_postprocessing(\n",
    "    input_image, segmentation_image,\n",
    "    smoothing_iterations, area_threshold, area_is_absolute,\n",
    "    use_sam_predictor, use_sam_autoseg, iou_threshold\n",
    "):\n",
    "    post_masks = {}\n",
    "    # if not use_sam_predictor and not use_sam_autoseg:\n",
    "    mask = postprocess(\n",
    "        segmentation_image, smoothing_iterations,\n",
    "        area_threshold, area_is_absolute\n",
    "    )\n",
    "    post_masks[\"Simple\"] = mask\n",
    "\n",
    "    if use_sam_predictor:\n",
    "        mask = postprocess_with_sam(\n",
    "            segmentation_image,\n",
    "            smoothing_iterations, area_threshold, area_is_absolute\n",
    "        )\n",
    "        post_masks[\"SAMPredictor\"] = mask\n",
    "\n",
    "    if use_sam_autoseg:\n",
    "        sam_auto_masks = get_sam_auto_masks(input_image)\n",
    "        mask = postprocess_with_sam_auto(\n",
    "            sam_auto_masks,\n",
    "            segmentation_image,\n",
    "            smoothing_iterations, iou_threshold,\n",
    "            area_threshold, area_is_absolute\n",
    "        )\n",
    "        post_masks[\"SAMAutoSegmentation\"] = mask\n",
    "\n",
    "\n",
    "    return post_masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set the Input, RF Model and the result directory paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input image\n",
    "data_path = \"../datasets/data.tif\"\n",
    "data_path = Path(data_path)\n",
    "print(f\"data_path exists: {data_path.exists()}\")\n",
    "\n",
    "# random forest model\n",
    "rf_model_path = \"../datasets/rf_model.bin\"\n",
    "rf_model_path = Path(rf_model_path)\n",
    "print(f\"rf_model_path exists: {rf_model_path.exists()}\")\n",
    "\n",
    "# result folder\n",
    "segmentation_dir = Path(\"../datasets/segmentation_result\")\n",
    "segmentation_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# temporary storage path for saving extracted embeddings patches\n",
    "storage_path = \"./temp_storage.hdf5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare the Input and RF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get patch sizes\n",
    "input_stack = Image.open(data_path)\n",
    "\n",
    "num_slices = input_stack.n_frames\n",
    "img_height = input_stack.height\n",
    "img_width = input_stack.width\n",
    "\n",
    "print(num_slices, img_height, img_width)\n",
    "# print(patch_size, target_patch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(rf_model_path, mode=\"rb\") as f:\n",
    "    model_data = pickle.load(f)\n",
    "# compatibility check for old format rf model\n",
    "if isinstance(model_data, dict):  # noqa: SIM108\n",
    "    # new format\n",
    "    rf_model = model_data[\"rf_model\"]\n",
    "else:\n",
    "    # old format\n",
    "    rf_model = model_data\n",
    "\n",
    "rf_model.set_params(verbose=0)\n",
    "rf_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing the Model for Feature Extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list of available models\n",
    "get_available_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"MobileSAM\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# print(f\"running on {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_adapter = get_model(model_name, img_height, img_width)\n",
    "\n",
    "patch_size = model_adapter.patch_size\n",
    "overlap = model_adapter.overlap\n",
    "\n",
    "patch_size, overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# post-processing parameters\n",
    "do_postprocess = True\n",
    "\n",
    "smoothing_iterations = 25\n",
    "area_threshold = 100        # to ignore mask regions with area below this threshold\n",
    "area_is_absolute = True    # is area is based on pixels or pecentage (False)\n",
    "\n",
    "use_sam_predictor = True\n",
    "use_sam_autoseg = False\n",
    "sam_autoseg_iou_threshold = 0.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the slice temporary storage\n",
    "storage = h5py.File(storage_path, \"w\")\n",
    "storage_group = storage.create_group(\"slice\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, page in tqdm(\n",
    "    enumerate(ImageSequence.Iterator(input_stack)),\n",
    "    desc=\"Slices\", total=num_slices, position=0\n",
    "):\n",
    "    # print(f\"slice {i + 1}\", end=\"\\n\")\n",
    "    slice_img = np.array(page.convert(\"RGB\"))\n",
    "\n",
    "    get_slice_features(slice_img, patch_size, overlap, model_adapter, storage_group)\n",
    "\n",
    "    segmentation_image = predict_slice(\n",
    "        rf_model, storage_group[model_adapter.name], model_adapter,\n",
    "        img_height, img_width,\n",
    "        patch_size, overlap\n",
    "    )\n",
    "\n",
    "    img = Image.fromarray(segmentation_image)\n",
    "    img.save(segmentation_dir.joinpath(f\"slice_{i:04}_prediction.tiff\"))\n",
    "\n",
    "    if do_postprocess:\n",
    "        post_masks = apply_postprocessing(\n",
    "            slice_img, segmentation_image,\n",
    "            smoothing_iterations, area_threshold, area_is_absolute,\n",
    "            use_sam_predictor, use_sam_autoseg, sam_autoseg_iou_threshold\n",
    "        )\n",
    "        # save results\n",
    "        for name, mask in post_masks.items():\n",
    "            img = Image.fromarray(mask)\n",
    "            seg_dir = segmentation_dir.joinpath(name)\n",
    "            seg_dir.mkdir(exist_ok=True) \n",
    "            img.save(seg_dir.joinpath(f\"slice_{i:04}_{name}.tiff\"))\n",
    "\n",
    "\n",
    "\n",
    "if storage is not None:\n",
    "    storage.close()\n",
    "    storage = None\n",
    "Path(storage_path).unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if storage is not None:\n",
    "    storage.close()\n",
    "    Path(storage_path).unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "project52",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}