DeePiCt_predict3d.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DeePiCt_predict3d.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOC87Ef9iQ4j7Y3eJ+sNZFa",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ZauggGroup/DeePiCt/blob/main/DeePiCt_predict3d.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This Colab notebook can be used to create predictions with already trained 3D models for ribosome, membrane, microtubule and FAS prediction. The tomogram that you want to use for prediction should be available somewhere online, for example on Drive. The initial step of the spectrum matching filter is not included in the notebook, and you should execute it beforehand. For more details about the model, follow the instructions available on the [DeePiCt Github repository](https://github.com/ZauggGroup/DeePiCt/blob/main/README.md).  \n",
        "\n",
        "## Instructions: \n",
        "* This notebook includes 4 steps to segment the tomogram and optional step 5 for visualization of the result. \n",
        "* Make sure that the tomogram with applied filter is available on an online share, example Google Drive. \n",
        "* Run the cells in the order that they are displayed. To run a cell, you need to click the play button on the left corner of the cell. \n",
        "* Some cells contain parameters that need to be defined, so make sure you enter all the required information correctly before running the particular cell. You have to run the cell, so that the parameter value is saved. \n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "6TVwiLkkBlcy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configurations\n",
        "___\n",
        "\n",
        "### Make sure you have GPU access enabled by going to Runtime -> Change Runtime Type -> Hardware accelerator and selecting GPU"
      ],
      "metadata": {
        "id": "n5e-5k5UCdlr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 1. Installations\n",
        "___"
      ],
      "metadata": {
        "id": "y9hZoeCxWv8q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 1.1. Run this cell to connect your Google Drive to colab\n",
        "\n",
        "#@markdown * Click on the URL. \n",
        "\n",
        "#@markdown * Sign in your Google Account. \n",
        "\n",
        "#@markdown You will either have to:\n",
        "#@markdown * copy the authorisation code and enter it into box below OR\n",
        "\n",
        "#@markdown * in the new google colab, you can just click \"Allow\" and it should connect.\n",
        "\n",
        "#@markdown * Click on \"Folder\" icon on the Left, press the refresh button. Your Google Drive folder should now be available here as \"gdrive\". \n",
        "\n",
        "# mount user's Google Drive to Google Colab.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "9DLz4_aOGuCo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4L8dPqnU-ehK",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown ## 1.2. Run this cell to install necessary packages\n",
        "\n",
        "#@markdown The code in this cell: \n",
        "#@markdown * Gets the git repository of DeePiCt\n",
        "\n",
        "!git clone https://github.com/ZauggGroup/DeePiCt.git\n",
        "\n",
        "#@markdown * Installs required packages\n",
        "\n",
        "!pip install mrcfile\n",
        "!pip install tensorboardX"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 2. Set the data variables and config file\n",
        "\n",
        "___"
      ],
      "metadata": {
        "id": "PFZauVVpXWae"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 2.1. Choose the model based on what you want to segment. The available models are prediction for membrane, ribosome, microtubules, and FAS. \n",
        "\n",
        "# Define the variable:\n",
        "predict_type = \"microtubule\" #@param [\"ribosome\",\"membrane\", \"microtubule\", \"FAS\"]\n",
        "\n",
        "models_weights = {\"ribosome\": \"https://www.dropbox.com/sh/oavbtcvusi07xbh/AAB4CRl9fY4hX9ypsAhwDrNna/3d_cnn/full_vpp_ribo_model_IF4_D2_BN.pth?dl=0\",\n",
        "                  \"membrane\": \"https://www.dropbox.com/sh/oavbtcvusi07xbh/AADm29QsXAHenTSTkASMcCk0a/3d_cnn/full_vpp_memb_model_IF4_D2_BN.pth?dl=0\", \n",
        "                  \"microtubule\": \"https://www.dropbox.com/sh/oavbtcvusi07xbh/AAAtAV19H3ar5nE2KYecqZf3a/3d_cnn/11cents5inter6mitotic_pp.pth?dl=0\", \n",
        "                  \"FAS\": \"https://www.dropbox.com/sh/oavbtcvusi07xbh/AACeAnRVDjCNy9viktZVdxiza/3d_cnn/full_vpp_fas_model_IF16_D2_BN.pth?dl=0\"}\n",
        "\n",
        "mapping = {'microtubule': 'microtubule', 'ribosome':'ribo', 'FAS':'fas', 'membrane':'mem'}\n",
        "class_name = mapping[predict_type]\n",
        "!wget -O model_weights.pth {models_weights[predict_type]}"
      ],
      "metadata": {
        "id": "OWqwPEkhJt2r",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from posixpath import split\n",
        "#@markdown ## 2.2. Define important variables\n",
        "\n",
        "#@markdown ### Define the following information in the given variables:\n",
        "\n",
        "scriptdir = '/content/DeePiCt/3d_cnn/scripts'\n",
        "srcdir = '/content/DeePiCt/3d_cnn/src/'\n",
        "original_config_file = '/content/DeePiCt/3d_cnn/config.yaml'\n",
        "model_path = '/content/model_weights.pth'\n",
        "\n",
        "# Define the folowing variables:\n",
        "\n",
        "# @markdown * **ID/name for the tomogram**:\n",
        "tomo_name = '180426_005' #@param {type:\"string\"}\n",
        "\n",
        "# @markdown * **Path to the tomogram .mrc file**:\n",
        "tomogram_path = '/content/gdrive/MyDrive/tomo_data/match_spectrum_filt.mrc' #@param {type:\"string\"}\n",
        "\n",
        "# @markdown * **Path to the mask .mrc file used for processing (if there is no mask leave it empty)**:\n",
        "mask_path = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ### You don't need to change the following variables:\n",
        "\n",
        "# @markdown * *Path where the config .yaml file will be saved (you can leave the default option)*:\n",
        "user_config_file = '/content/gdrive/MyDrive/DeePiCt_3d/config.yaml'  #@param {type:\"string\"}\n",
        "\n",
        "# @markdown * *Path where the data .csv file will be saved (you can leave the default option)*:\n",
        "user_data_file = '/content/gdrive/MyDrive/DeePiCt_3d/data.csv' #@param {type:\"string\"}\n",
        "\n",
        "# @markdown * *Path to folder where the prediction files will be saved (you can leave the default option)*:\n",
        "user_prediction_folder = '/content/gdrive/MyDrive/DeePiCt_3d/'  #@param {type:\"string\"}\n",
        "\n",
        "# @markdown *  *Path to folder where the intermediate files will be saved (you can leave the default option)*:\n",
        "user_work_folder = '/content/work/'  #@param {type:\"string\"}\n",
        "\n",
        "import os\n",
        "\n",
        "os.makedirs(os.path.split(user_config_file)[0], exist_ok=True)\n",
        "os.makedirs(os.path.split(user_data_file)[0], exist_ok=True)\n",
        "os.makedirs(os.path.split(user_prediction_folder)[0], exist_ok=True)\n",
        "os.makedirs(os.path.split(user_work_folder)[0], exist_ok=True)\n"
      ],
      "metadata": {
        "id": "5S-BJrqshNTy",	
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 2.3. Create data csv file and yaml config file\n",
        "#@markdown Run this cell to create the .csv data file and .yaml config file\n",
        "\n",
        "import csv\n",
        "import yaml\n",
        "\n",
        "header = ['tomo_name','raw_tomo','filtered_tomo', 'no_mask']\n",
        "\n",
        "# Define the elements of this list:\n",
        "data = [tomo_name, '', tomogram_path, mask_path]\n",
        "\n",
        "with open(user_data_file, 'w', encoding='UTF8') as f:\n",
        "    writer = csv.writer(f)\n",
        "\n",
        "    # write the header\n",
        "    writer.writerow(header)\n",
        "\n",
        "    # write the data\n",
        "    writer.writerow(data)\n",
        "  \n",
        "data_dictionary = dict(zip(header, data))\n",
        "\n",
        "def read_yaml(file_path):\n",
        "    with open(file_path, \"r\") as stream:\n",
        "        data = yaml.safe_load(stream)\n",
        "    return data\n",
        "\n",
        "def save_yaml(data, file_path):\n",
        "    with open(file_path, 'w') as yaml_file:\n",
        "        yaml.dump(data, yaml_file, default_flow_style=False)\n",
        "\n",
        "d = read_yaml(original_config_file)\n",
        "d['dataset_table'] = user_data_file\n",
        "d['output_dir'] = user_prediction_folder\n",
        "d['work_dir'] = user_work_folder\n",
        "d['model_path'] = f'{model_path}'\n",
        "d['tomos_sets']['training_list'] = []\n",
        "d['tomos_sets']['prediction_list'] = [f'{tomo_name}']\n",
        "d['cross_validation']['active'] = False\n",
        "d['training']['active'] = False\n",
        "d['prediction']['active'] = True\n",
        "d['evaluation']['particle_picking']['active'] = False\n",
        "d['evaluation']['segmentation_evaluation']['active'] = False\n",
        "d['training']['processing_tomo'] = 'filtered_tomo'\n",
        "d['training']['semantic_classes'] = [class_name]\n",
        "d['prediction']['processing_tomo'] = 'filtered_tomo'\n",
        "d['prediction']['semantic_class'] = class_name\n",
        "d['postprocessing_clustering']['region_mask'] = 'no_mask'\n",
        "save_yaml(d, user_config_file)"
      ],
      "metadata": {
        "id": "7HYB2B9-CMHL",	
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 3. Predict with trained neural network\n",
        "\n",
        "___"
      ],
      "metadata": {
        "id": "Hmk0sQ2r_DR5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 3.1. Split tomogram into patches\n",
        "#@markdown Run this cell to create the 64^3 patches\n",
        "\n",
        "!python /content/DeePiCt/3d_cnn/scripts/generate_prediction_partition.py --config_file {user_config_file} --pythonpath '/content/DeePiCt/3d_cnn/src' --tomo_name {tomo_name}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "kAO1qyqaBnop"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 3.2. Segment the patches\n",
        "#@markdown Run this cell to create the segmentation of the 64^3 patches \n",
        "\n",
        "#@markdown *Note*: If you have disk full warning, try to ignore it and let the cell continue to run\n",
        "\n",
        "!python /content/DeePiCt/3d_cnn/scripts/segment.py --config_file {user_config_file} --pythonpath '/content/DeePiCt/3d_cnn/src' --tomo_name {tomo_name} --gpu 0"
      ],
      "metadata": {
        "id": "SpNSS0Th_HHa",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 3.3. Assemble the patches\n",
        "#@markdown Run this cell to assemble the segmented patches\n",
        "\n",
        "!python /content/DeePiCt/3d_cnn/scripts/assemble_prediction.py --config_file {user_config_file} --pythonpath '/content/DeePiCt/3d_cnn/src' --tomo_name {tomo_name}"
      ],
      "metadata": {
        "id": "ceHyX9jX_KIR",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 4. Post-processing of the prediction\n",
        "\n",
        "___\n",
        "\n"
      ],
      "metadata": {
        "id": "Dn9-434P_KXw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 4.1. Post-processing of the prediction\n",
        "#@markdown Run this cell to do post-processing of the predictions\n",
        "\n",
        "import yaml\n",
        "\n",
        "def read_yaml(file_path):\n",
        "    with open(file_path, \"r\") as stream:\n",
        "        data = yaml.safe_load(stream)\n",
        "    return data\n",
        "\n",
        "def save_yaml(data, file_path):\n",
        "    with open(file_path, 'w') as yaml_file:\n",
        "        yaml.dump(data, yaml_file, default_flow_style=False)\n",
        "\n",
        "d = read_yaml(user_config_file)\n",
        "\n",
        "#@markdown #### If you don't want to use the default parameters, unclick the button for `default_options` and define the parameters. Otherwise, the default options will be used.\n",
        "\n",
        "default_options = True #@param {type:\"boolean\"}\n",
        "\n",
        "if default_options:\n",
        "    d['postprocessing_clustering']['active'] = True\n",
        "    d['postprocessing_clustering']['threshold'] = 0.5\n",
        "    d['postprocessing_clustering']['min_cluster_size'] = 500\n",
        "    d['postprocessing_clustering']['max_cluster_size'] = None\n",
        "    d['postprocessing_clustering']['clustering_connectivity'] = 1\n",
        "    d['postprocessing_clustering']['calculate_motl'] = True\n",
        "    d['postprocessing_clustering']['ignore_border_thickness'] = 0\n",
        "    d['postprocessing_clustering']['region_mask'] = 'no_mask'\n",
        "    d['postprocessing_clustering']['contact_mode'] = 'intersection'\n",
        "    d['postprocessing_clustering']['contact_distance'] = 0\n",
        "else:\n",
        "    threshold = 0.5 #@param {type:\"number\"}\n",
        "    min_cluster_size = 500 #@param {type:\"integer\"}\n",
        "    max_cluster_size = 0 #@param {type:\"integer\"}\n",
        "    clustering_connectivity = 1 #@param {type:\"integer\"}\n",
        "    calculate_motl = True #@param {type:\"boolean\"}\n",
        "    contact_mode = 'intersection' #@param [\"contact\", \"colocalization\", \"intersection\"]\n",
        "    contact_distance = 0 #@param {type:\"integer\"}\n",
        "    if max_cluster_size == 0:\n",
        "        max_cluster_size = None\n",
        "    d['postprocessing_clustering']['active'] = True\n",
        "    d['postprocessing_clustering']['threshold'] = threshold\n",
        "    d['postprocessing_clustering']['min_cluster_size'] = min_cluster_size\n",
        "    d['postprocessing_clustering']['max_cluster_size'] = max_cluster_size\n",
        "    d['postprocessing_clustering']['clustering_connectivity'] = clustering_connectivity\n",
        "    d['postprocessing_clustering']['calculate_motl'] = calculate_motl\n",
        "    d['postprocessing_clustering']['ignore_border_thickness'] = 0\n",
        "    d['postprocessing_clustering']['region_mask'] = 'no_mask'\n",
        "    d['postprocessing_clustering']['contact_mode'] = 'intersection'\n",
        "    d['postprocessing_clustering']['contact_distance'] = contact_distance\n",
        "\n",
        "save_yaml(d, user_config_file)\n",
        "\n",
        "!python /content/DeePiCt/3d_cnn/scripts/clustering_and_cleaning.py --config_file {user_config_file} --pythonpath '/content/DeePiCt/3d_cnn/src' --tomo_name {tomo_name}"
      ],
      "metadata": {
        "id": "OzjLaSbp_QDH",	
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ***NOTE:*** if the previous cell is failing because of memory error, try this slower version. Otherwise, the post-processing is done and you can skip this cell.\n",
        "\n",
        "!python /content/DeePiCt/3d_cnn/scripts/clustering_and_cleaning_slower.py --config_file {user_config_file} --pythonpath '/content/DeePiCt/3d_cnn/src' --tomo_name {tomo_name}"
      ],
      "metadata": {
        "id": "Ty41RhnPrTag",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 5. Visualize results\n",
        "\n",
        "___"
      ],
      "metadata": {
        "id": "vqUgAoaBsjtb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 5.1. Read the tomogram and the prediction\n",
        "#@markdown Run this cell to read the tomogram and the predictions\n",
        "\n",
        "import mrcfile\n",
        "\n",
        "\n",
        "def read_tomogram(filename):\n",
        "    with mrcfile.open(filename, permissive=True) as m:\n",
        "        return m.data\n",
        "\n",
        "tomogram = read_tomogram(data_dictionary['filtered_tomo'])\n",
        "path = f\"{d['output_dir']}/predictions/{d['model_path'].split('/')}\"\n",
        "predictions = read_tomogram(f'{user_prediction_folder}/predictions/model_weights/{tomo_name}/{class_name}/post_processed_prediction.mrc')"
      ],
      "metadata": {
        "id": "CCYMplv6skCD",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## 5.2. Visualize the prediction\n",
        "#@markdown Run this cell to do visualization of particular z axis\n",
        "\n",
        "\n",
        "z_axis = 200 #@param {type:\"integer\"}\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# First figure\n",
        "plt.figure(figsize = (10,10))\n",
        "plt.imshow(tomogram[z_axis], cmap='gray')\n",
        "\n",
        "# Second figure\n",
        "plt.figure(figsize = (10,10))\n",
        "plt.imshow(tomogram[z_axis], cmap='gray')\n",
        "alphas = np.zeros(predictions.shape)\n",
        "alphas[predictions > 0] = 1\n",
        "plt.imshow(predictions[z_axis], alpha=alphas[z_axis], cmap='Blues')"
      ],
      "metadata": {
        "id": "7R1IMHON2FKJ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}