Embeddings/create_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "<div>\n",
    "<img src=\"https://github.com/fsemerar/SlicerTomoSAM/raw/main/TomoSAM/Resources/Media/tomosam_logo.png\" width=\"300\"/>\n",
    "</div>"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "This notebook helps with the generation of the image embeddings for all the slices of your tiff stack along the three Cartesian directions. You can create the embeddings by running this notebook either locally or on Colab. A GPU is recommended for this step to speed up the process; in Colab, make sure to select `Runtime`→`Change runtime type` and set the `Hardware accelerator` to GPU. Locally, you will first need to create the conda environment, as shown in the README."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2viNExExM4up",
    "outputId": "8533fecb-0cb1-41be-ea34-858337e01b8a",
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "from google.colab import output\n",
    "output.enable_custom_widget_manager()\n",
    "!pip install https://github.com/facebookresearch/segment-anything/archive/6fdee8f2727f4506cfbbe553e23b895e27956588.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download weights for SAM\n",
    "![ ! -f \"sam_vit_h_4b8939.pth\" ] && wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vxl6cH4eNXM8"
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "import torch\n",
    "import sys, os\n",
    "import pickle\n",
    "def create_embeddings(img_input_filepath, output_filepath, sam_checkpoint_path):\n",
    "\n",
    "    check, img = cv2.imreadmulti(img_input_filepath)\n",
    "    img = np.array(img)\n",
    "    if not check:\n",
    "        raise Exception(\"Image file not found.\")\n",
    "    elif img.ndim > 3 or img.ndim < 2:\n",
    "        raise Exception(\"Unsupported image type.\")\n",
    "    elif img.ndim == 2:\n",
    "        img = img[:, :, np.newaxis]\n",
    "\n",
    "    print(f\"Image dimensions: {img.shape}\")\n",
    "\n",
    "    sam = sam_model_registry[\"vit_h\"](checkpoint=sam_checkpoint_path)\n",
    "    if torch.cuda.is_available():\n",
    "        sam.to(device=\"cuda\")\n",
    "    predictor = SamPredictor(sam)\n",
    "\n",
    "    embeddings = [[], [], []]\n",
    "    slice_direction = ['x', 'y', 'z']\n",
    "    for i, d in enumerate(slice_direction):\n",
    "        print(f\"\\nSlicing along {d} direction\")\n",
    "        for k in range(img.shape[i]):\n",
    "            if i == 0:\n",
    "                img_slice = img[k]\n",
    "            elif i == 1:\n",
    "                img_slice = img[:, k]\n",
    "            else:\n",
    "                img_slice = img[:, :, k]\n",
    "            sys.stdout.write(f\"\\rCreating embedding for {k + 1}/{img.shape[i]} image\")\n",
    "            predictor.reset_image()\n",
    "            if torch.cuda.is_available():\n",
    "                torch.cuda.empty_cache()\n",
    "            predictor.set_image(np.repeat(img_slice[:, :, np.newaxis], 3, axis=2))\n",
    "            embeddings[i].append({'original_size': predictor.original_size,\n",
    "                                  'input_size': predictor.input_size,\n",
    "                                  'features': predictor.features.to('cpu')})\n",
    "\n",
    "    with open(output_filepath + \".pkl\", 'wb') as f:\n",
    "        pickle.dump(embeddings, f)\n",
    "        print(f\"\\nSaved {output_filepath}.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 73
    },
    "id": "R6ollXf-NnuH",
    "outputId": "e3495438-9d70-43d1-adc2-fd3ddc69fb5a"
   },
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "img_filename = list(files.upload().keys())[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for local use\n",
    "img_filename = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2O-YSz09OdXp",
    "outputId": "73bb0f52-7b05-4b68-93dc-ab0a410f6796"
   },
   "outputs": [],
   "source": [
    "create_embeddings(img_filename, os.path.splitext(img_filename)[0], \"sam_vit_h_4b8939.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4_IWf18DPZHw"
   },
   "outputs": [],
   "source": [
    "# Download from Colab\n",
    "files.download(img_filename + \".pkl\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}