cellpose_2D_TD_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca1c2304-a831-487c-b83e-54b812ac36e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q cellpose\n",
    "!pip install -q ipywidgets\n",
    "!pip install -q matplotlib --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "92532e3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU is available!\n"
     ]
    }
   ],
   "source": [
    "from cellpose import utils, io, models, plot\n",
    "from ipywidgets import interact, interact_manual\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from tifffile import imread, imwrite\n",
    "import time\n",
    "import sys\n",
    "import os, random\n",
    "import torch\n",
    "if torch.cuda.is_available():\n",
    "    print(\"GPU is available!\")\n",
    "else:\n",
    "    print(\"GPU is NOT available! Are you on the right node?\")\n",
    "    exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9ec0e545",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logging.basicConfig(\n",
    "    format='%(asctime)s [%(levelname)s] %(name)s - %(message)s',\n",
    "    level=logging.INFO,\n",
    "    datefmt='%Y-%m-%d %H:%M:%S',\n",
    "    stream=sys.stdout,\n",
    ")\n",
    "log = logging.getLogger()\n",
    "fhandler = logging.FileHandler(filename='cellpose_nuclei_log.txt', mode='w')\n",
    "formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s - %(message)s')\n",
    "fhandler.setFormatter(formatter)\n",
    "log.addHandler(fhandler)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cab8d9db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-03-22 23:41:08 [INFO] cellpose.core - ** TORCH CUDA version installed and working. **\n",
      "2023-03-22 23:41:08 [INFO] cellpose.core - >>>> using GPU\n",
      "2023-03-22 23:41:08 [INFO] cellpose.models - >> nuclei << model set to be used\n",
      "2023-03-22 23:41:08 [INFO] cellpose.models - >>>> model diam_mean =  17.000 (ROIs rescaled to this size during training)\n",
      "Nuclei model enabled\n",
      "--------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc142abd096a4ae992f826417b3b1ec3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Dropdown(description='file', options=('20230308_AWV_F05_t106_c003_poseinput.tiff', '2023…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81268b2660c94cbe9d1281c54f4ad8c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Button(description='Run Interact', style=ButtonStyle()), Output()), _dom_classes=('widge…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.batch_process()>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#@markdown ### Provide the path to your dataset and to the folder where the predictions are saved, then play the cell to predict outputs from your unseen images.\n",
    "\n",
    "Data_folder = \"/research/rgs01/home/clusterHome/tdas/cellpose_input\" #@param {type:\"string\"}\n",
    "Result_folder = \"/research/rgs01/home/clusterHome/tdas/cellpose_nuclei\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown ###Are your data single images or stacks?\n",
    "\n",
    "Data_type = \"Single_Images\" #@param [\"Single_Images\", \"Stacks (2D + t)\"]\n",
    "\n",
    "#@markdown ###What model do you want to use?\n",
    "\n",
    "model_choice = \"Nuclei\" #@param [\"Cytoplasm\",\"Cytoplasm2\", \"Cytoplasm2_Omnipose\", \"Bacteria_Omnipose\", \"Nuclei\", \"Own_model\"]\n",
    "\n",
    "#@markdown ####If using your own model, please provide the path to the model (not the folder):\n",
    "\n",
    "Prediction_model = \"\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown ### What channel do you want to segment?\n",
    "\n",
    "Channel_to_segment= \"Blue\" #@param [\"Grayscale\", \"Blue\", \"Green\", \"Red\"]\n",
    "\n",
    "# @markdown ###If you chose the model \"cytoplasm\" indicate if you also have a nuclear channel that can be used to aid the segmentation.\n",
    "\n",
    "Nuclear_channel= \"Blue\" #@param [\"None\", \"Blue\", \"Green\", \"Red\"]\n",
    "\n",
    "#@markdown ### Segmentation parameters:\n",
    "Object_diameter =  150 #@param {type:\"number\"}\n",
    "\n",
    "Flow_threshold = 0.4 #@param {type:\"slider\", min:0.1, max:1.1, step:0.1}\n",
    "\n",
    "#cellprob_threshold = 0 #@param {type:\"slider\", min:-6, max:6, step:1}\n",
    "\n",
    "\n",
    "# Find the number of channel in the input image\n",
    "\n",
    "random_choice = random.choice(os.listdir(Data_folder))\n",
    "x = io.imread(Data_folder+\"/\"+random_choice)\n",
    "n_channel = 1 if x.ndim == 2 else x.shape[-1]\n",
    "\n",
    "if Channel_to_segment == \"Grayscale\":\n",
    "  segment_channel = 0\n",
    "\n",
    "  if Data_type == \"Single_Images\":\n",
    "    if not n_channel == 1:\n",
    "        print(bcolors.WARNING +\"!! WARNING: your image has more than one channel, choose which channel you want to use for your predictions !!\")\n",
    "\n",
    "if Channel_to_segment == \"Blue\":\n",
    "  segment_channel = 3\n",
    "\n",
    "if Channel_to_segment == \"Green\":\n",
    "  segment_channel = 2\n",
    "\n",
    "if Channel_to_segment == \"Red\":\n",
    "  segment_channel = 1\n",
    "\n",
    "if Nuclear_channel == \"Blue\":\n",
    "  nuclear_channel = 3\n",
    "\n",
    "if Nuclear_channel == \"Green\":\n",
    "  nuclear_channel = 2\n",
    "\n",
    "if Nuclear_channel == \"Red\":\n",
    "  nuclear_channel = 1\n",
    "\n",
    "if Nuclear_channel == \"None\":\n",
    "  nuclear_channel = 0\n",
    "\n",
    "if model_choice == \"Cytoplasm\":  \n",
    "  channels=[segment_channel,nuclear_channel]\n",
    "  model = models.Cellpose(gpu=True, model_type=\"cyto\")\n",
    "  print(\"Cytoplasm model enabled\")\n",
    "\n",
    "if model_choice == \"Cytoplasm2\":  \n",
    "  channels=[segment_channel,nuclear_channel]\n",
    "  model = models.Cellpose(gpu=True, model_type=\"cyto2\")\n",
    "  print(\"Cytoplasm2 model enabled\")\n",
    "\n",
    "if model_choice == \"Cytoplasm2_Omnipose\":  \n",
    "  channels=[segment_channel,nuclear_channel]\n",
    "  model = models.Cellpose(gpu=True, model_type=\"cyto2_omni\")\n",
    "  print(\"Cytoplasm2_Omnipose model enabled\")\n",
    "  \n",
    "if model_choice == \"Nuclei\":\n",
    "  channels=[segment_channel,0]\n",
    "  model = models.Cellpose(gpu=True, model_type=\"nuclei\")\n",
    "  print(\"Nuclei model enabled\")\n",
    "\n",
    "if model_choice == \"Bacteria_Omnipose\":\n",
    "  channels=[segment_channel,nuclear_channel]\n",
    "  model = models.Cellpose(gpu=True, model_type=\"bact_omni\")\n",
    "  Object_diameter =  0\n",
    "  print(\"Bacteria_omnipose model enabled\")\n",
    "\n",
    "if model_choice == \"Own_model\":\n",
    "  channels=[segment_channel,nuclear_channel]\n",
    "  model = models.CellposeModel(gpu=True, pretrained_model=Prediction_model, torch=True, diam_mean=30.0, net_avg=True, device=None, residual_on=True, style_on=True, concatenation=False)\n",
    "\n",
    "  print(\"Own model enabled\")\n",
    "\n",
    "if Object_diameter == 0:\n",
    "  Object_diameter = None\n",
    "  print(\"The cell size will be estimated automatically for each image\")\n",
    "\n",
    "if Data_type == \"Single_Images\" :\n",
    "\n",
    "  print('--------------------------------------------------------------')\n",
    "  @interact\n",
    "  def preview_results(file = os.listdir(Data_folder)):\n",
    "    source_image = io.imread(os.path.join(Data_folder, file))\n",
    "    \n",
    "    if model_choice == \"Own_model\":\n",
    "      masks, flows, styles = model.eval(source_image, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "    \n",
    "    else:\n",
    "      masks, flows, styles, diams = model.eval(source_image, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "    \n",
    "    flowi = flows[0]\n",
    "    fig = plt.figure(figsize=(20,10))\n",
    "    plot.show_segmentation(fig, source_image, masks, flowi, channels=channels)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "  def batch_process():\n",
    "      print(\"Your images are now beeing processed\")\n",
    "      for name in os.listdir(Data_folder):\n",
    "        print(\"Performing prediction on: \"+name)\n",
    "        image = io.imread(Data_folder+\"/\"+name)\n",
    "        short_name = os.path.splitext(name)\n",
    "\n",
    "        if model_choice == \"Own_model\":\n",
    "          masks, flows, styles = model.eval(image, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "        else:\n",
    "          masks, flows, styles, diams = model.eval(image, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "\n",
    "        os.chdir(Result_folder)\n",
    "        imwrite(str(short_name[0])+\"_nuclearmask.tif\", masks)\n",
    "\n",
    "  im = interact_manual(batch_process)\n",
    "  im.widget.children[0].description = 'Process your images'\n",
    "  im.widget.children[0].style.button_color = 'yellow'\n",
    "  display(im)\n",
    "\n",
    "if Data_type == \"Stacks (2D + t)\" :\n",
    "  print(\"Stacks (2D + t) are now beeing predicted\")\n",
    "  \n",
    "  print('--------------------------------------------------------------')\n",
    "  @interact\n",
    "  def preview_results_stacks(file = os.listdir(Data_folder)):\n",
    "    timelapse = imread(Data_folder+\"/\"+file)\n",
    "\n",
    "    if model_choice == \"Own_model\":\n",
    "      masks, flows, styles = model.eval(timelapse[0], diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "    else:\n",
    "      masks, flows, styles, diams = model.eval(timelapse[0], diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "           \n",
    "    flowi = flows[0]\n",
    "    fig = plt.figure(figsize=(20,10))\n",
    "    plot.show_segmentation(fig, timelapse[0], masks, flowi, channels=channels)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "  def batch_process_stack():\n",
    "      print(\"Your images are now beeing processed\")  \n",
    "      for image in os.listdir(Data_folder):\n",
    "        print(\"Performing prediction on: \"+image)\n",
    "        timelapse = imread(Data_folder+\"/\"+image)\n",
    "        short_name = os.path.splitext(image)\n",
    "        n_timepoint = timelapse.shape[0]\n",
    "        prediction_stack = np.zeros((n_timepoint, timelapse.shape[1], timelapse.shape[2]))\n",
    "        \n",
    "        for t in range(n_timepoint):\n",
    "          print(\"Frame number: \"+str(t))\n",
    "          img_t = timelapse[t]\n",
    "\n",
    "          if model_choice == \"Own_model\":\n",
    "            masks, flows, styles = model.eval(img_t, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "          else:\n",
    "            masks, flows, styles, diams = model.eval(img_t, diameter=Object_diameter, flow_threshold=Flow_threshold, channels=channels)\n",
    "            \n",
    "              \n",
    "          prediction_stack[t] = masks\n",
    "      \n",
    "        prediction_stack_32 = img_as_float32(prediction_stack, force_copy=False)\n",
    "        os.chdir(Result_folder)\n",
    "        imwrite(str(short_name[0])+\".tif\", prediction_stack_32)\n",
    "  \n",
    "  im = interact_manual(batch_process_stack)\n",
    "  im.widget.children[0].description = 'Process your images'\n",
    "  im.widget.children[0].style.button_color = 'yellow'\n",
    "  display(im)          \n",
    "     \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6e1075f-9d7c-4370-b789-4d541f587d54",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:hpcf-jupyter-py3.9]",
   "language": "python",
   "name": "conda-env-hpcf-jupyter-py3.9-py"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}