TreeLearn_Pipeline.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zZ0XVazsDLEM"
      },
      "source": [
        "## Example Notebook for TreeLearn\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KTNLHg1A4hTJ"
      },
      "source": [
        "Thank you for your interest in our TreeLearn method! With this notebook and google colab, you can try out the pipeline for segmenting a forest point cloud without installing anything on your own computer! However, to actually use our method, we recommend to set up the environment on a gpu-capable device and run the segmentation pipeline there, as described in our repository.\n",
        "\n",
        "You need to be signed in with your google account. Please also make sure that you are connected to a gpu runtime by by selecting 'runtime' change runtime to e.g. T4 GPU. The following code snippet will show a table with gpu information if you are connnected to a gpu runtime. To run the code snippet, simply click on the left edge. or press (Ctrl + enter) after selecting it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DI42xeR55Qsl",
        "outputId": "f020121a-7745-48a9-90f7-f536f1d5c4c6"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fCl6z0uNseOV"
      },
      "source": [
        "The following two code snippets are necessary to set up the environment and download the model weights. Simply run them before continuing. It takes 2 to 3 minutes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "_SYchH0wLXqx"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "# install environment\n",
        "!pip install torch-scatter -f https://data.pyg.org/whl/torch-2.1.0+cu118.html\n",
        "!pip install timm==0.6.12\n",
        "!pip install tensorboard\n",
        "!pip install tensorboardX\n",
        "!pip install laspy[lazrs]==2.5.1\n",
        "!pip install munch==2.5.0\n",
        "!pip install pandas==2.0.0\n",
        "!pip install plyfile==0.9\n",
        "!pip install pyyaml==6.0\n",
        "!pip install scikit-learn==1.2.2\n",
        "!pip install six==1.16.0\n",
        "!pip install tqdm==4.65.0\n",
        "!pip install open3d-cpu==0.17.0 --default-timeout=100\n",
        "!pip install jakteristics==0.5.1\n",
        "!pip install shapely==2.0.1\n",
        "!pip install geopandas==0.12.2\n",
        "!pip install alphashape==1.3.1\n",
        "!pip install spconv-cu114 --default-timeout=100"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "0qxJEvNu9kGm"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!git clone https://github.com/ecker-lab/TreeLearn.git\n",
        "%cd TreeLearn\n",
        "!pip install -e .\n",
        "%cd ..\n",
        "!mkdir data\n",
        "!mkdir checkpoints\n",
        "!mkdir pipeline\n",
        "!mkdir pipeline/forests\n",
        "path = \"/content/checkpoints/model_weights.pth\"\n",
        "!wget https://data.goettingen-research-online.de/api/access/datafile/:persistentId?persistentId=doi:10.25625/VPMPID/1JMEQV -O $path\n",
        "\n",
        "%cd TreeLearn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ra8mJVagRvlI"
      },
      "source": [
        "## Pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "INqgac0r0KQJ"
      },
      "source": [
        "We first need to decide which point cloud we want to segment. The following code snippet downloads an example point cloud segment that we did not train on.\n",
        "If you want to try out another forest point cloud, replace the download with your own. Make sure that the file is in the .npy, .npz, .las, .laz or .txt file format and the total size of the forest plot should be around 1600 m^2 at maximum. Please note that with a forest point cloud of this size the segmentation took in our runs around 15 minutes in google colab due to limited computation resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "25148d-10JxJ",
        "outputId": "b9f66503-1a21-48f7-8c85-a8a5853a2e10"
      },
      "outputs": [],
      "source": [
        "forest_name = \"plot_7_cut.laz\"\n",
        "path = \"/content/pipeline/forests/\" + forest_name\n",
        "!wget https://data.goettingen-research-online.de/api/access/datafile/:persistentId?persistentId=doi:10.25625/VPMPID/0WDXL6 -O $path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZREpAa-s0Fp"
      },
      "source": [
        "\n",
        "To run the TreeLearn pipeline interactively in google colab, we import the function run_treelearn_pipeline. This function takes as argument the config dict. \n",
        "We import the pipeline.yaml as the config dict and print it. We adjust some entries in the config dict to fit to the setting in google colab and speed up the pipeline. We also initialize the logger so that the progress in the pipeline is printed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 711
        },
        "id": "bPbQGPxNdwzL",
        "outputId": "2add54bd-b356-423d-dc06-87121e565761"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.path.append(\"/content/TreeLearn/tools/pipeline\")\n",
        "from pipeline import run_treelearn_pipeline\n",
        "import argparse, pprint\n",
        "from tree_learn.util import get_config\n",
        "\n",
        "config_path = \"/content/TreeLearn/configs/pipeline/pipeline.yaml\"\n",
        "config = get_config(config_path)\n",
        "\n",
        "# adjust config\n",
        "config.forest_path = \"/content/pipeline/forests/\" + forest_name\n",
        "config.dataset_test.data_root = \"/content/pipeline/tiles\"\n",
        "config.tile_generation = True\n",
        "config.pretrain = \"/content/checkpoints/model_weights.pth\"\n",
        "config.sample_generation.stride = 0.9 # small overlap of tiles\n",
        "config.shape_cfg.outer_remove = False # default value = 13.5\n",
        "config.save_cfg.save_treewise = False\n",
        "config.save_cfg.return_type = \"voxelized_and_filtered\"\n",
        "print(pprint.pformat(config.toDict(), indent=2))\n",
        "\n",
        "import logging\n",
        "logger = logging.getLogger(\"TreeLearn\")\n",
        "for handler in logger.handlers[:]:\n",
        "    logger.removeHandler(handler)\n",
        "logging.basicConfig()\n",
        "ch = logging.StreamHandler(sys.stdout)\n",
        "formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n",
        "ch.setFormatter(formatter)\n",
        "logger.addHandler(ch)\n",
        "logger.setLevel(logging.INFO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6neoLf6l2zcF"
      },
      "source": [
        "After having set all the correct settings in the config file, it remains to run the pipeline. Please keep in mind that fully running it for the example point cloud takes around 12 minutes.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JKh7yZ3ld1o3"
      },
      "outputs": [],
      "source": [
        "# run pipeline\n",
        "run_treelearn_pipeline(config)\n",
        "# runtime ~ 12 min"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zroEoOMU5DOO"
      },
      "source": [
        "If everything has run as expected, the segmented point cloud is now saved in the .laz format with labels in the directory /content/pipeline/results. It is also saved in the .npz format. You can easily download it by right-clicking and selecting download. **Please note that, to speed up the pipeline for demonstration purposes, we do not remove a buffer of 13.5 meters at the edge of the input point cloud. This is usually recommended since points at the edge do not have the necessary context for making accurate predictions. The predictions at the edge of the example point cloud do not represent the capabilities of TreeLearn.**"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}