notebooks/s3dis_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# S3DIS inference\n",
    "This notebook lets you run inference on the S3DIS Test Fold, from a pretrained model's checkpoint on your machine. The output will be a confusion matrix which you can save locally on your machine. For this, you must have done the following:\n",
    "- downloaded and preprocessed the S3DIS dataset (if you haven't, instantiating the dataset will launch it for you though)\n",
    "- trained a model (which produced a checkpoint directory) or downloaded our pretrained weights\n",
    "\n",
    "Note this notebook should work for multimodal (2D+3D) and unimodal (3D only) experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select you GPU\n",
    "I_GPU = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to use autoreload\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "import glob\n",
    "from omegaconf import OmegaConf\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "torch.cuda.set_device(I_GPU)\n",
    "DIR = os.path.dirname(os.getcwd())\n",
    "ROOT = os.path.join(DIR, \"..\")\n",
    "sys.path.insert(0, ROOT)\n",
    "sys.path.insert(0, DIR)\n",
    "\n",
    "from torch_points3d.utils.config import hydra_read\n",
    "from torch_points3d.trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set your parameters\n",
    "DATA_ROOT = '/path/to/your/dataset/root/directory'\n",
    "checkpoint_dir = '/directory/containing/your/checkpoint/files'        # each fold's checkpoint should: follow checkpoint_dir/area_i/model_name.pt\n",
    "result_dir = '/directory/where/to/save/inference/metrics'\n",
    "model_name = 'Res16UNet34-L4-early'                                   # adapt if you use another model in your checkpoint\n",
    "sample_res = 1                                                        # controls the sphere sampling grid resolution. The lower, the more val/test spheres and the higher their overlap\n",
    "n_votes = 1                                                           # number of inferences per spherical sample. For multi-inference voting with inference-time augmentation\n",
    "batch_size = 8                                                        # increase if your device allows it\n",
    "full_res = True                                                       # predictions will be made on the raw point cloud, at full resolution\n",
    "num_workers = 4                                                       # increase if your machine allows it\n",
    "exp_name = None                                                       # you may give a name to the experiment\n",
    "exp_name = f'{model_name}_votes-{n_votes}_sample_res-{sample_res}' if exp_name is None else exp_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Inference on S3DIS exp={exp_name}')\n",
    "\n",
    "for i_area in range(6):\n",
    "    print(f'\\nInference on Area {i_area + 1}')\n",
    "    \n",
    "    # These are the arguments passed to Hydra. You could run the same thing \n",
    "    # from CLI using `eval.py` with the command `python eval.py [...]`\n",
    "    overrides = [\n",
    "        f'model_name={model_name}',\n",
    "        f\"checkpoint_dir={os.path.join(checkpoint_dir, f'area_{i + 1}')}\",\n",
    "        f'voting_runs={n_votes}',\n",
    "        f'tracker_options.full_res={full_res}',\n",
    "        'precompute_multi_scale=False',\n",
    "        f'num_workers={num_workers}',\n",
    "        f'batch_size={batch_size}',\n",
    "        f'cuda={I_GPU}',\n",
    "        'weight_name=latest',\n",
    "        f'+data.dataroot={DATA_ROOT}',\n",
    "        f'+data.eval_sample_res={eval_sample_res}',\n",
    "    ]\n",
    "\n",
    "    # Parse the arguments with Hydra and OmegaConf\n",
    "    cfg = hydra_read(overrides, config_name='eval')\n",
    "    OmegaConf.set_struct(cfg, False)\n",
    "\n",
    "    # Create the Trainer instance from your checkpoint\n",
    "    trainer = Trainer(cfg)\n",
    "\n",
    "    # Update the test transforms to match train transforms for test-time \n",
    "    # augmentation\n",
    "    trainer._dataset.test_dataset[0].transform = trainer._dataset.train_dataset.transform\n",
    "    if trainer._model.is_multimodal:\n",
    "        trainer._dataset.test_dataset[0].transform_image = trainer._dataset.train_dataset.transform_image\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[4].use_coverage = False\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[4].credit = int(2097152 * 2)\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[5].sigma = 0.02\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[6].transform.saturation = [0.8, 1.2]\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[6].transform.brightness = [0.8, 1.2]\n",
    "        trainer._dataset.test_dataset[0].transform_image.transforms[6].transform.contrast = [0.8, 1.2]\n",
    "\n",
    "    # Run inference\n",
    "    trainer.eval(stage_name='test')\n",
    "    cm = trainer._tracker._full_confusion_matrix\n",
    "    torch.save(cm, f'{result_dir}/{exp_name}_area-{i_area + 1}.pt')\n",
    "    print(f'  mIoU={cm.get_average_intersection_union() * 100:0.2f}')\n",
    "    print(f'  OA={cm.get_overall_accuracy() * 100:0.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confusion matrices will be saved in your `result_dir`.\n",
    "\n",
    "Now, if we want to compute the 6-Fold performance, we need to accumulate the confusion matrices from all the folds into a single one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm_6fold = None\n",
    "\n",
    "for i_area in range(6):\n",
    "    \n",
    "    # Load the confusion matrix previously computed\n",
    "    path = f'{result_dir}/{exp_name}_area-{i_area + 1}.pt'\n",
    "    cm_fold = torch.load(path)\n",
    "    \n",
    "    # Simple case for the first fold\n",
    "    if i_area == 0:\n",
    "        cm_6fold = cm_fold\n",
    "        continue\n",
    "    \n",
    "    # Accumulate the CMs from all folds\n",
    "    cm_6fold.confusion_matrix += cm_fold.confusion_matrix\n",
    "\n",
    "    # Save the 6-fold CM\n",
    "torch.save(cm_6fold, f'{result_dir}/{exp_name}_6-fold.pt')\n",
    "    \n",
    "print(f'6-Fold inference')\n",
    "print(f'  mIoU={cm_6fold.get_average_intersection_union() * 100:0.2f}')\n",
    "print(f'  OA={cm_6fold.get_overall_accuracy() * 100:0.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please refer to `torch_points3d/metrics/confusion_matrix` if you need to compute more metrics from the confusion matrices."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tp3d_dev]",
   "language": "python",
   "name": "conda-env-tp3d_dev-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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}