notebooks/4_pbr_segmentation_kpconv.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6710e0b",
   "metadata": {},
   "source": [
    "## todo: 1. kpconv training; 2. inference; 3. inference result visualization; 4. documentation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b44a8d7",
   "metadata": {},
   "source": [
    "## Create dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "377291ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add package path\n",
    "import sys\n",
    "import os\n",
    "path = os.path.abspath(os.path.join(os.getcwd(), os.pardir))\n",
    "if path not in sys.path:\n",
    "    sys.path.insert(0, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "72d44e8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dreamslab_legion/.local/lib/python3.7/site-packages/numba/core/errors.py:149: UserWarning: Insufficiently recent colorama version found. Numba requires colorama >= 0.3.9\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "# import packages\n",
    "from omegaconf import OmegaConf\n",
    "import pyvista as pv\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm.auto import tqdm\n",
    "import time\n",
    "import laspy\n",
    "\n",
    "from rock_detection_3d.datasets.segmentation.rock_las import RockLASDataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd6ee2d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# configure visualization\n",
    "os.environ[\"DISPLAY\"] = \":1.0\"\n",
    "os.environ[\"PYVISTA_OFF_SCREEN\"]=\"true\"\n",
    "os.environ[\"PYVISTA_PLOT_THEME\"]=\"true\"\n",
    "os.environ[\"PYVISTA_USE_PANEL\"]=\"true\"\n",
    "os.environ[\"PYVISTA_AUTO_CLOSE\"]=\"false\"\n",
    "os.system(\"Xvfb :1 -screen 0 1024x768x24 > /dev/null 2>&1 &\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "673c0616",
   "metadata": {},
   "outputs": [],
   "source": [
    "# configure dataset params\n",
    "\n",
    "DIR = \"\" # Replace with your root directory, the data will go in DIR/data.\n",
    "USE_COLOR = True #@param {type:\"boolean\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e66b5cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "pbr_yaml = \"\"\"\n",
    "class: None # shapenet.ShapeNetDataset\n",
    "task: segmentation\n",
    "dataroot: %s\n",
    "color: %r                                     # Use color vectors as features\n",
    "first_subsampling: 0.02                       # Grid size of the input data\n",
    "pre_transforms:                               # Offline transforms, done only once        \n",
    "    - transform: GridSampling3D\n",
    "      params:\n",
    "        size: ${first_subsampling}\n",
    "train_transforms:                             # Data augmentation pipeline\n",
    "    - transform: RandomNoise\n",
    "      params:\n",
    "        sigma: 0.01\n",
    "        clip: 0.05\n",
    "    - transform: RandomScaleAnisotropic\n",
    "      params:\n",
    "        scales: [0.9,1.1]\n",
    "    - transform: AddOnes\n",
    "    - transform: AddFeatsByKeys\n",
    "      params:\n",
    "        list_add_to_x: [True]\n",
    "        feat_names: [\"ones\"]\n",
    "        delete_feats: [True]\n",
    "test_transforms:\n",
    "    - transform: AddOnes\n",
    "    - transform: AddFeatsByKeys\n",
    "      params:\n",
    "        list_add_to_x: [True]\n",
    "        feat_names: [\"ones\"]\n",
    "        delete_feats: [True]\n",
    "\"\"\" % (os.path.join(DIR,\"data\"), USE_COLOR) \n",
    "\n",
    "from omegaconf import OmegaConf\n",
    "params = OmegaConf.create(pbr_yaml)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eb208749",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset: RockLASDataset \n",
       "\u001b[0;95mtrain_pre_batch_collate_transform \u001b[0m= None\n",
       "\u001b[0;95mval_pre_batch_collate_transform \u001b[0m= None\n",
       "\u001b[0;95mtest_pre_batch_collate_transform \u001b[0m= None\n",
       "\u001b[0;95mpre_transform \u001b[0m= Compose([\n",
       "    GridSampling3D(grid_size=0.02, quantize_coords=False, mode=mean),\n",
       "])\n",
       "\u001b[0;95mtest_transform \u001b[0m= Compose([\n",
       "    AddOnes(),\n",
       "    AddFeatsByKeys(ones=True),\n",
       "])\n",
       "\u001b[0;95mtrain_transform \u001b[0m= Compose([\n",
       "    RandomNoise(sigma=0.01, clip=0.05),\n",
       "    RandomScaleAnisotropic([0.9, 1.1]),\n",
       "    AddOnes(),\n",
       "    AddFeatsByKeys(ones=True),\n",
       "])\n",
       "\u001b[0;95mval_transform \u001b[0m= None\n",
       "\u001b[0;95minference_transform \u001b[0m= Compose([\n",
       "    GridSampling3D(grid_size=0.02, quantize_coords=False, mode=mean),\n",
       "    AddOnes(),\n",
       "    AddFeatsByKeys(ones=True),\n",
       "])\n",
       "Size of \u001b[0;95mtrain_dataset \u001b[0m= 29\n",
       "Size of \u001b[0;95mtest_dataset \u001b[0m= 10\n",
       "Size of \u001b[0;95mval_dataset \u001b[0m= 10\n",
       "\u001b[0;95mBatch size =\u001b[0m None"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create dataset\n",
    "dataset = RockLASDataset(params)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a78e53b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "640705679fef440f9a72cda36b56bb7a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ViewInteractiveWidget(height=412, layout=Layout(height='auto', width='100%'), width=1024)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visually inspect dataset \n",
    "\n",
    "#@title Plot samples with part annotations { run: \"auto\" }\n",
    "objectid_1 = 3 #@param {type:\"slider\", min:0, max:100, step:1}\n",
    "objectid_2 = 4 #@param {type:\"slider\", min:0, max:100, step:1}\n",
    "objectid_3 = 5 #@param {type:\"slider\", min:0, max:100, step:1}\n",
    "\n",
    "samples = [objectid_1,objectid_2,objectid_3]\n",
    "p = pv.Plotter(notebook=True,shape=(1, len(samples)),window_size=[1024,412])\n",
    "for i in range(len(samples)):\n",
    "    p.subplot(0, i)\n",
    "    sample = dataset.train_dataset[samples[i]]\n",
    "    point_cloud = pv.PolyData(sample.pos.numpy())\n",
    "    point_cloud['y'] = sample.y.numpy()\n",
    "    p.add_points(point_cloud,  show_scalar_bar=False, point_size=4)\n",
    "    p.camera_position = [-1,5, -10]\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4976467",
   "metadata": {},
   "source": [
    "## Create segmentation model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3cb00924",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages\n",
    "from torch_points3d.applications.kpconv import KPConv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "517d23f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SegKPConv(\n",
       "  (unet): KPConvUnet(\n",
       "    (down_modules): ModuleList(\n",
       "      (0): KPDualBlock(\n",
       "        Nb parameters: 34304\n",
       "        (blocks): ModuleList(\n",
       "          (0): SimpleBlock(\n",
       "            Nb parameters: 3968; None; RadiusNeighbourFinder {'_radius': 0.05, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "            (kp_conv): KPConvLayer(InF: 4, OutF: 64, kernel_pts: 15, radius: 0.03, KP_influence: linear, Add_one: False)\n",
       "            (bn): FastBatchNorm1d(\n",
       "              (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "          (1): ResnetBBlock(\n",
       "            Nb parameters: 30336\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 15424; None; RadiusNeighbourFinder {'_radius': 0.05, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 32, OutF: 32, kernel_pts: 15, radius: 0.03, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(32, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=64, out_features=32, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(32, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=32, out_features=128, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Sequential(\n",
       "              (0): Linear(in_features=64, out_features=128, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1): KPDualBlock(\n",
       "        Nb parameters: 144000\n",
       "        (blocks): ModuleList(\n",
       "          (0): ResnetBBlock(\n",
       "            Nb parameters: 23936\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 15424; GridSampling3D(grid_size=0.04, quantize_coords=False, mode=mean); RadiusNeighbourFinder {'_radius': 0.05, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 32, OutF: 32, kernel_pts: 15, radius: 0.03, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(32, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=32, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(32, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=32, out_features=128, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Identity()\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "          (1): ResnetBBlock(\n",
       "            Nb parameters: 120064\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 61568; None; RadiusNeighbourFinder {'_radius': 0.1, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 64, OutF: 64, kernel_pts: 15, radius: 0.06, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=64, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=64, out_features=256, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=256, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (2): KPDualBlock(\n",
       "        Nb parameters: 572672\n",
       "        (blocks): ModuleList(\n",
       "          (0): ResnetBBlock(\n",
       "            Nb parameters: 94976\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 61568; GridSampling3D(grid_size=0.08, quantize_coords=False, mode=mean); RadiusNeighbourFinder {'_radius': 0.1, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 64, OutF: 64, kernel_pts: 15, radius: 0.06, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=256, out_features=64, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=64, out_features=256, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Identity()\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "          (1): ResnetBBlock(\n",
       "            Nb parameters: 477696\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 246016; None; RadiusNeighbourFinder {'_radius': 0.2, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 128, OutF: 128, kernel_pts: 15, radius: 0.12, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=256, out_features=128, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Sequential(\n",
       "              (0): Linear(in_features=256, out_features=512, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (3): KPDualBlock(\n",
       "        Nb parameters: 2284032\n",
       "        (blocks): ModuleList(\n",
       "          (0): ResnetBBlock(\n",
       "            Nb parameters: 378368\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 246016; GridSampling3D(grid_size=0.16, quantize_coords=False, mode=mean); RadiusNeighbourFinder {'_radius': 0.2, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 128, OutF: 128, kernel_pts: 15, radius: 0.12, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=128, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Identity()\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "          (1): ResnetBBlock(\n",
       "            Nb parameters: 1905664\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 983552; None; RadiusNeighbourFinder {'_radius': 0.4, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 256, OutF: 256, kernel_pts: 15, radius: 0.24, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=256, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=256, out_features=1024, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(1024, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(1024, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (4): KPDualBlock(\n",
       "        Nb parameters: 9122816\n",
       "        (blocks): ModuleList(\n",
       "          (0): ResnetBBlock(\n",
       "            Nb parameters: 1510400\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 983552; GridSampling3D(grid_size=0.32, quantize_coords=False, mode=mean); RadiusNeighbourFinder {'_radius': 0.4, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 256, OutF: 256, kernel_pts: 15, radius: 0.24, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=1024, out_features=256, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=256, out_features=1024, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(1024, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Identity()\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "          (1): ResnetBBlock(\n",
       "            Nb parameters: 7612416\n",
       "            (kp_conv): SimpleBlock(\n",
       "              Nb parameters: 3933184; None; RadiusNeighbourFinder {'_radius': 0.8, '_max_num_neighbors': 25, '_conv_type': 'partial_dense'}\n",
       "              (kp_conv): KPConvLayer(InF: 512, OutF: 512, kernel_pts: 15, radius: 0.48, KP_influence: linear, Add_one: False)\n",
       "              (bn): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (activation): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_1): Sequential(\n",
       "              (0): Linear(in_features=1024, out_features=512, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (unary_2): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=2048, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(2048, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "              (2): LeakyReLU(negative_slope=0.1)\n",
       "            )\n",
       "            (shortcut_op): Sequential(\n",
       "              (0): Linear(in_features=1024, out_features=2048, bias=False)\n",
       "              (1): FastBatchNorm1d(\n",
       "                (batch_norm): BatchNorm1d(2048, eps=1e-05, momentum=0.02, affine=True, track_running_stats=True)\n",
       "              )\n",
       "            )\n",
       "            (activation): LeakyReLU(negative_slope=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (inner_modules): ModuleList(\n",
       "      (0): Identity()\n",
       "    )\n",
       "    (up_modules): ModuleList(\n",
       "      (0): FPModule_PD(\n",
       "        (nn): Sequential(\n",
       "          (0): Sequential(\n",
       "            (0): Linear(in_features=3072, out_features=512, bias=False)\n",
       "            (1): FastBatchNorm1d(\n",
       "              (batch_norm): BatchNorm1d(512, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): LeakyReLU(negative_slope=0.2)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1): FPModule_PD(\n",
       "        (nn): Sequential(\n",
       "          (0): Sequential(\n",
       "            (0): Linear(in_features=1024, out_features=256, bias=False)\n",
       "            (1): FastBatchNorm1d(\n",
       "              (batch_norm): BatchNorm1d(256, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): LeakyReLU(negative_slope=0.2)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (2): FPModule_PD(\n",
       "        (nn): Sequential(\n",
       "          (0): Sequential(\n",
       "            (0): Linear(in_features=512, out_features=128, bias=False)\n",
       "            (1): FastBatchNorm1d(\n",
       "              (batch_norm): BatchNorm1d(128, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): LeakyReLU(negative_slope=0.2)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (3): FPModule_PD(\n",
       "        (nn): Sequential(\n",
       "          (0): Sequential(\n",
       "            (0): Linear(in_features=256, out_features=64, bias=False)\n",
       "            (1): FastBatchNorm1d(\n",
       "              (batch_norm): BatchNorm1d(64, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): LeakyReLU(negative_slope=0.2)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (mlp): Sequential(\n",
       "      (0): Sequential(\n",
       "        (0): Linear(in_features=64, out_features=2, bias=False)\n",
       "        (1): FastBatchNorm1d(\n",
       "          (batch_norm): BatchNorm1d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): LeakyReLU(negative_slope=0.2)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create KPConv model\n",
    "\n",
    "color = 3  # use RGB data\n",
    "\n",
    "class SegKPConv(torch.nn.Module):\n",
    "    def __init__(self, cat_to_seg):\n",
    "        super().__init__()\n",
    "        self.unet = KPConv(\n",
    "            architecture = \"unet\", \n",
    "            input_nc = color, \n",
    "            output_nc = 2,  # isPBR & notPBR\n",
    "            num_layers= 4, \n",
    "            in_grid_size = params['first_subsampling'],  # grid size at the entry of the network; should be consistent of dataset first sampling resolution\n",
    "            )\n",
    "    \n",
    "    @property\n",
    "    def conv_type(self):\n",
    "        \"\"\" This is needed by the dataset to infer which batch collate should be used\"\"\"\n",
    "        return self.unet.conv_type\n",
    "    \n",
    "    def get_batch(self):\n",
    "        return self.batch\n",
    "    \n",
    "    def get_output(self):\n",
    "        \"\"\" This is needed by the tracker to get access to the ouputs of the network\"\"\"\n",
    "        return self.output\n",
    "    \n",
    "    def get_labels(self):\n",
    "        \"\"\" Needed by the tracker in order to access ground truth labels\"\"\"\n",
    "        return self.labels\n",
    "    \n",
    "    \n",
    "    def get_current_losses(self):\n",
    "        \"\"\" Entry point for the tracker to grab the loss \"\"\"\n",
    "        return {\"loss_seg\": float(self.loss_seg)}\n",
    "\n",
    "    def forward(self, data):\n",
    "        self.labels = data.y\n",
    "        self.batch = data.batch\n",
    "        \n",
    "        # Forward through unet and classifier\n",
    "        output_batch = self.unet(data)\n",
    "        \n",
    "        self.output = output_batch.x\n",
    "        #print(self.output)\n",
    "        #print(self.labels)\n",
    "        \n",
    "        # Set loss for the backward pass\n",
    "        self.loss_seg = torch.nn.functional.nll_loss(self.output, self.labels)\n",
    "        return self.output\n",
    "\n",
    "    def get_spatial_ops(self):\n",
    "        return self.unet.get_spatial_ops()\n",
    "        \n",
    "    def backward(self):\n",
    "         self.loss_seg.backward() \n",
    "\n",
    "\n",
    "model = SegKPConv(dataset.class_to_segments)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe2da58e",
   "metadata": {},
   "source": [
    "## The data loaders and CPU pre computing features\n",
    "KPConv is quite demanding on spatial operations such as grid sampling and radius search. On the network loaded here we have 10 KPConv layers on the encoder which means 10 radius search operations with varying number of neighbours. We observed a significant performance gain by moving those operations to the CPU where they can easily be optimised with suitable data structures such as kd-tree. We use [nonaflann](https://github.com/jlblancoc/nanoflann) in the back-end, a 3D optimised kd-tree implementation. Note that this is beneficiary only if you have access to multiple CPU threads.\n",
    "\n",
    "You can decide to precompute those spatial operations by setting the `precompute_multi_scale` parameter to `True` when creating the data loaders. The dataset will mine the model to figure out which spatial operations are required and in which order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1065eb3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_WORKERS = 2  # 4\n",
    "BATCH_SIZE = 2  # 16\n",
    "dataset.create_dataloaders(\n",
    "    model,\n",
    "    batch_size=BATCH_SIZE, \n",
    "    num_workers=NUM_WORKERS, \n",
    "    shuffle=True, \n",
    "    precompute_multi_scale=True \n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "37a9c438",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['x',\n",
       " 'y',\n",
       " 'pos',\n",
       " 'multiscale',\n",
       " 'upsample',\n",
       " 'batch',\n",
       " 'category',\n",
       " 'center',\n",
       " 'file_name',\n",
       " 'grid_size',\n",
       " 'id_scan',\n",
       " 'origin_id',\n",
       " 'ptr',\n",
       " 'scale']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = next(iter(dataset.train_dataloader))\n",
    "sample.keys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57780a10",
   "metadata": {},
   "source": [
    "Our `sample` contains the pre computed spatial information in the `multiscale` (encoder side) and `upsample` (decoder) attrivutes. The decoder pre computing is quite simple and just involves some basic caching for the nearest neighbour interpolation operation. Let's take a look at the encoder side of things first. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4e2ed34a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Batch(batch=[17482], idx_neighboors=[17482, 25], pos=[17482, 3]),\n",
       " Batch(batch=[17482], idx_neighboors=[17482, 25], pos=[17482, 3]),\n",
       " Batch(batch=[6995], grid_size=[2], idx_neighboors=[6995, 25], pos=[6995, 3]),\n",
       " Batch(batch=[6995], idx_neighboors=[6995, 25], pos=[6995, 3]),\n",
       " Batch(batch=[1706], grid_size=[2], idx_neighboors=[1706, 25], pos=[1706, 3]),\n",
       " Batch(batch=[1706], idx_neighboors=[1706, 25], pos=[1706, 3]),\n",
       " Batch(batch=[404], grid_size=[2], idx_neighboors=[404, 25], pos=[404, 3]),\n",
       " Batch(batch=[404], idx_neighboors=[404, 25], pos=[404, 3]),\n",
       " Batch(batch=[116], grid_size=[2], idx_neighboors=[116, 25], pos=[116, 3]),\n",
       " Batch(batch=[116], idx_neighboors=[116, 25], pos=[116, 3])]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample.multiscale"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "262a80c9",
   "metadata": {},
   "source": [
    "`sample.multiscale` contains 10 different versions of the input batch, each one of these versions contains the location of the points in `pos` as well as the neighbours of these points in the previous point cloud. We will first look at the points coming out of each downsampling layer (strided convolution), we have 5 of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ce098cd2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac04c7c263c04006b5fd75542393853e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ViewInteractiveWidget(height=256, layout=Layout(height='auto', width='100%'), width=1024)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#@title Successive downsampling {run:\"auto\"}\n",
    "sample_in_batch = 0 #@param {type:\"slider\", min:0, max:5, step:1}\n",
    "ms_data = sample.multiscale \n",
    "num_downsize = int(len(ms_data) / 2)\n",
    "p = pv.Plotter(notebook=True,shape=(1, num_downsize),window_size=[1024,256])\n",
    "for i in range(0,num_downsize):\n",
    "    p.subplot(0, i)\n",
    "    pos = ms_data[2*i].pos[ms_data[2*i].batch == sample_in_batch].numpy()\n",
    "    point_cloud = pv.PolyData(pos)\n",
    "    point_cloud['y'] = pos[:,1]\n",
    "    p.add_points(point_cloud,  show_scalar_bar=False, point_size=3)\n",
    "    p.add_text(\"Layer {}\".format(i+1),font_size=10)\n",
    "    p.camera_position = [-1,5, -10]\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc53cbcc",
   "metadata": {},
   "source": [
    "Let's now take one point in a layer (query point) and show its neighbours in the previous layer (support point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1e77a56e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cea73ce7a4cb41789c3c95d5cbf13d85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ViewInteractiveWidget(height=412, layout=Layout(height='auto', width='100%'), width=1024)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#@title Explore Neighborhood {run: \"auto\"}\n",
    "selected_layer = 7 #@param {type:\"slider\", min:1, max:9, step:1}\n",
    "sample_in_batch = 0 #@param {type:\"slider\", min:0, max:5, step:1}\n",
    "point1_id = 3 #@param {type:\"slider\", min:0, max:600, step:1}\n",
    "point2_id =  8 #@param {type:\"slider\", min:0, max:600, step:1}\n",
    "\n",
    "p = pv.Plotter(notebook=True,shape=(1, 2),window_size=[1024,412])\n",
    "\n",
    "# Selected layer\n",
    "p.subplot(0, 1)\n",
    "ms_data = sample.multiscale[selected_layer]\n",
    "pos = ms_data.pos[ms_data.batch == sample_in_batch].numpy()\n",
    "nei = ms_data.idx_neighboors[ms_data.batch == sample_in_batch]\n",
    "point_cloud = pv.PolyData(pos)\n",
    "p.add_points(point_cloud,  show_scalar_bar=False, point_size=3,opacity=0.3)\n",
    "p.add_points(pos[point1_id,:],  show_scalar_bar=False, point_size=7.0,color='red')\n",
    "p.add_points(pos[point2_id,:],  show_scalar_bar=False, point_size=7.0,color='green')\n",
    "p.camera_position = [-1,5, -10]\n",
    "\n",
    "# Previous layer\n",
    "p.subplot(0, 0)\n",
    "ms_data = sample.multiscale[selected_layer-1]\n",
    "pos = ms_data.pos[ms_data.batch == sample_in_batch].numpy()\n",
    "point_cloud = pv.PolyData(pos)\n",
    "p.add_points(point_cloud,  show_scalar_bar=False,point_size=3, opacity=0.3)\n",
    "nei_pos = ms_data.pos[nei[point1_id]].numpy()\n",
    "nei_pos = nei_pos[nei[point1_id] >= 0]\n",
    "p.add_points(nei_pos,  show_scalar_bar=False, point_size=3.0,color='red')\n",
    "nei_pos = ms_data.pos[nei[point2_id]].numpy()\n",
    "nei_pos = nei_pos[nei[point2_id] >= 0]\n",
    "p.add_points(nei_pos,  show_scalar_bar=False, point_size=3.0,color='green')\n",
    "p.camera_position = [-1,5, -10]\n",
    "\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18bc6d9c",
   "metadata": {},
   "source": [
    "## Train neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d1544081",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trainer:\n",
    "    def __init__(self, model, dataset, num_epoch = 60, device=torch.device('cuda'), checkpoint_path=\"model/kpconv\"):\n",
    "        self.num_epoch = num_epoch\n",
    "        self._model = model\n",
    "        self._dataset=dataset\n",
    "        self.device = device\n",
    "        self.checkpoint_path = checkpoint_path\n",
    "        if not os.path.exists(self.checkpoint_path):\n",
    "            os.makedirs(self.checkpoint_path)\n",
    "        \n",
    "    def save_model(self, epoch):\n",
    "        f = os.path.join(self.checkpoint_path, \"{epoch}.pt\".format(epoch=epoch))\n",
    "        torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': self._model.state_dict(),\n",
    "            'optimizer_state_dict': self.optimizer.state_dict(),\n",
    "            'loss': self._model.loss_seg,\n",
    "            }, f)\n",
    "        \n",
    "    def load_model(self, epoch):\n",
    "        f = os.path.join(self.checkpoint_path, \"{epoch}.pt\".format(epoch=epoch))\n",
    "        assert os.path.isfile(f)\n",
    "        checkpoint = torch.load(f)\n",
    "        self._model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        self._model.to(self.device)\n",
    "        self._model.eval()\n",
    "        return self._model\n",
    "        \n",
    "    def resume_model(self, epoch):\n",
    "        f = os.path.join(self.checkpoint_path, \"{epoch}.pt\".format(epoch=epoch))\n",
    "        assert os.path.isfile(f)\n",
    "        self.optimizer = torch.optim.Adam(self._model.parameters(), lr=0.001)\n",
    "        self.tracker = self._dataset.get_tracker(False, True)\n",
    "        \n",
    "        checkpoint = torch.load(f)\n",
    "        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        self._model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        self._model.to(self.device)\n",
    "        \n",
    "\n",
    "    def fit(self):\n",
    "        self.optimizer = torch.optim.Adam(self._model.parameters(), lr=0.001)\n",
    "        self.tracker = self._dataset.get_tracker(False, True)\n",
    "\n",
    "        for i in range(self.num_epoch):\n",
    "            print(\"=========== EPOCH %i ===========\" % i)\n",
    "            time.sleep(0.5)\n",
    "            self.train_epoch()\n",
    "            self.tracker.publish(i)\n",
    "            self.valid_epoch()\n",
    "            self.tracker.publish(i)\n",
    "            self.save_model(i)\n",
    "\n",
    "    def train_epoch(self):\n",
    "        self._model.to(self.device)\n",
    "        self._model.train()\n",
    "        self.tracker.reset(\"train\")\n",
    "        train_loader = self._dataset.train_dataloader\n",
    "        iter_data_time = time.time()\n",
    "        with tqdm(train_loader) as tq_train_loader:\n",
    "            for i, data in enumerate(tq_train_loader):\n",
    "                t_data = time.time() - iter_data_time\n",
    "                iter_start_time = time.time()\n",
    "                self.optimizer.zero_grad()\n",
    "                data.to(self.device)\n",
    "                self._model.forward(data)\n",
    "                self._model.backward()\n",
    "                self.optimizer.step()\n",
    "                if i % 10 == 0:\n",
    "                    self.tracker.track(self._model)\n",
    "\n",
    "                tq_train_loader.set_postfix(\n",
    "                    **self.tracker.get_metrics(),\n",
    "                    data_loading=float(t_data),\n",
    "                    iteration=float(time.time() - iter_start_time),\n",
    "                )\n",
    "                iter_data_time = time.time()\n",
    "\n",
    "    def valid_epoch(self):\n",
    "        self._model.to(self.device)\n",
    "        self._model.eval()\n",
    "        self.tracker.reset(\"val\")\n",
    "        val_loader = self._dataset.val_dataloader\n",
    "        iter_data_time = time.time()\n",
    "        with tqdm(val_loader) as tq_val_loader:\n",
    "            for i, data in enumerate(tq_val_loader):\n",
    "                t_data = time.time() - iter_data_time\n",
    "                iter_start_time = time.time()\n",
    "                data.to(self.device)\n",
    "                self._model.forward(data)           \n",
    "                self.tracker.track(self._model)\n",
    "\n",
    "                tq_val_loader.set_postfix(\n",
    "                    **self.tracker.get_metrics(),\n",
    "                    data_loading=float(t_data),\n",
    "                    iteration=float(time.time() - iter_start_time),\n",
    "                )\n",
    "                iter_data_time = time.time()\n",
    "                \n",
    "    def test(self, save=True):\n",
    "        self._model.to(self.device)\n",
    "        self._model.eval()\n",
    "        self.tracker.reset(\"test\")\n",
    "        test_loader = self._dataset.test_dataloaders[0]\n",
    "        iter_data_time = time.time()\n",
    "        with tqdm(test_loader) as tq_test_loader:\n",
    "            for i, data in enumerate(tq_test_loader):\n",
    "                t_data = time.time() - iter_data_time\n",
    "                iter_start_time = time.time()\n",
    "                data.to(self.device)\n",
    "                self._model.forward(data)           \n",
    "                self.tracker.track(self._model)\n",
    "\n",
    "                tq_test_loader.set_postfix(\n",
    "                    **self.tracker.get_metrics(),\n",
    "                    data_loading=float(t_data),\n",
    "                    iteration=float(time.time() - iter_start_time),\n",
    "                )\n",
    "                if save:\n",
    "                    self.save_las(data)\n",
    "                \n",
    "                iter_data_time = time.time()\n",
    "                \n",
    "                \n",
    "        self.tracker.publish(self.num_epoch - 1)\n",
    "        if save:\n",
    "            print(\"saved prediction las files\")\n",
    "        \n",
    "         \n",
    "        \n",
    "    def save_las(self, batch_data):\n",
    "        pred = self._model.get_output()\n",
    "        outputs = torch.argmax(pred, 1)\n",
    "\n",
    "        for i in torch.unique(batch_data.batch):\n",
    "            idx = batch_data.batch==i\n",
    "            pos = batch_data.pos[idx]\n",
    "            x = batch_data.x[idx]\n",
    "            y = outputs[idx]\n",
    "            origin_id = batch_data.origin_id[idx]\n",
    "            id_scan = batch_data.id_scan[i]\n",
    "            grid_size = batch_data.grid_size[i]\n",
    "            category = batch_data.category[idx]\n",
    "            center = batch_data.center[i*3:i*3+3]\n",
    "            scale = batch_data.scale[i]\n",
    "            file_name = batch_data.file_name[i]\n",
    "            \n",
    "            self.write_las(file_name, pos, x, y, center, scale)\n",
    "\n",
    "            \n",
    "    def write_las(self, file_name, pos, x, y, center, scale):\n",
    "        path = 'data/rocklas/prediction_3d'\n",
    "        if not os.path.exists('data/rocklas/prediction_3d'):\n",
    "            os.makedirs(self.checkpoint_path)\n",
    "        pos = pos / scale\n",
    "        pos = pos + center\n",
    "        pos = pos.cpu().detach().numpy()\n",
    "        color = (x * (2**16)).cpu().detach().numpy().astype(np.uint16)\n",
    "        PBR_ids = (y==1).cpu().detach().numpy()\n",
    "        notPBR_ids = (y==0).cpu().detach().numpy()\n",
    "        isPBR = np.empty(pos.shape[0])\n",
    "        isPBR[:] = np.NaN\n",
    "        notPBR = isPBR.copy()\n",
    "        isPBR[PBR_ids] = 0\n",
    "        notPBR[notPBR_ids] = 1\n",
    "\n",
    "        f = os.path.join(path, \"pred_{f}\".format(f=file_name))\n",
    "        header = laspy.LasHeader(point_format=2, version=\"1.2\")\n",
    "        header.scales = np.array([0.01, 0.01, 0.01])\n",
    "        header.add_extra_dim(laspy.ExtraBytesParams(name=\"isPBR\", type=np.float64))\n",
    "        header.add_extra_dim(laspy.ExtraBytesParams(name=\"notPBR\", type=np.float64))\n",
    "\n",
    "        las = laspy.LasData(header)\n",
    "        las.x = pos[:, 0]\n",
    "        las.y = pos[:, 1]\n",
    "        las.z = pos[:, 2]\n",
    "        las.red = color[:, 0]\n",
    "        las.green =  color[:, 1]\n",
    "        las.blue = color[:, 2]\n",
    "        las.isPBR = isPBR\n",
    "        las.notPBR = notPBR\n",
    "\n",
    "        las.write(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d9ed2b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(model, dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4cc58542",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 0 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df40995f9c2d427586dfa883645f670e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e351db809bea470ebcfb10f0a23b533b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 1 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e7fed52ce0f4ca7a7969f1219cb09ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c73a8d4a04194a0bab581bef6566ce7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 2 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "24c2770b79f54f888ed90f04603c0fb6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed42406c2c524735b7f104f3fe1faa16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 3 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "58c8184afcf1405c932fe0580cd8a2e7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b73ba4da70f7433f8cf22b9315438051",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 4 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5e4ef3a5ec7477c9a55685f06ad9f90",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "204a3cc6b10e4f1f82c1ca1f1c3a0894",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 5 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcec4bddd1b246909e564dad495c451a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98de169a80c343dd913795649497282d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 6 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e1bb9fd0239482cbc383e41c4ef834c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f0a05c5d64342598327213ca539e559",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 7 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c53b3051633f4de5925a2d0f2f04fcd1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ff4b513014e44c7882306f32923f5800",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 8 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "897b8df45dfe4043a97fda6c5d2eb99b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e40a57ab506647b8993476d07819be3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 9 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1f9133663c04e0686e9172d3eb915ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9e1dd62e4ab4e65aa11cb0c6c8c0e82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 10 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d863860ca6a04f218b668c91a3e44ac5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "08c7092f80e34686984efd0a5cc6a688",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 11 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e2109de77a3430697fa17fa6dd292c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a18132ed24fd4f11899b142eb64cc47c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 12 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c50c9008d2134bf68a070430d2b81d64",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d6ad20b642c4ba196fda97ae02199d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 13 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9645b2f72a5a4aefaf504f93ca3700fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4ac42076f8b4410963650a3fabe3564",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 14 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bad998b8945f4d4a950e605796548daa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d11e7fe65484a0599a6106ffdd77480",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 15 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a4dd266b0d1d4b5f981410af111730eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a56d1c06ebdb464098537e470f660e89",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 16 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5fa2a9486a1485399b2c6d70108803b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "06791f7d91f8405fb7a92a23e2dfa0f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 17 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "547a9092107f477fa9146be022ca3d2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "897df7b54eec4fbea368b35f664f6ea9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 18 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b64c02b195c64184804b4ec4f8e0f4c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e59506ac32440d9a83595f1c3e2cc12",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 19 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0213c70e73704c33b6135de132362776",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8eb0619419864ffc93fce0c696c07de9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 20 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4feeb32021464602a631b859b45544d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3df8cd1cd90b472293a82bf13267f24f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 21 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c778c4439ea840bdaae89146d2a29e9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "215785e1bdef4ba9afc4e11a518a0e43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 22 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb354b484e6847fca66e2ffd74d29369",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc4df4b30d07416a8c534c5d6241706b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 23 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1af66b8220cc4467865417204aa32515",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a464af30c2045e0b6d1da02ddc52376",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 24 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80cf63e7027a48b2a40a886e36fd916d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "820bce81594841eb938e2afd3b4a2519",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 25 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a8f36a17cd04e5094fef6bc8f951799",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a3b2426b9684c6a83d511caa982ccd8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 26 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f05d94acb30d4cd8b9fecbe55c7c4320",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b6b425a04b8412fbb6f30d414748172",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 27 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "514b29c9235346288528dd7ce53f6482",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a84ea15f23dc47338ac7649127aca641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 28 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b879b4239c7c4f9ab865c21c44ececcc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c40d2e3357e4483bab3345970b9d3a14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 29 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6b1f350b3264b0197747ac03b5722b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a7e80d5a87c940238e2c894fe884d296",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 30 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "983c8b3299304a3a9f00fe2d4967e3b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12932f61ac7545898215ef5f5b8aee2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 31 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1cce72f89bb4d6a893c7e5bd19514c8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63f8ca4a6cdc47eeb4cf5ce84da9145b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 32 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca1fd15f634a4d099cc5915fbead618c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "885cfe3b4c9f4313acf97b54917cc44e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 33 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "343f02b05cd644d39662be4da65c079c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "357856b043b2498e852f59dbc8c5fed3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 34 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16f7187da5ab46ffa3e8ddcf1a5d798f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6b3640e53b7423790c3e323a73b5673",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 35 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e1eb62720f846309d1fb9b6f686ba20",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "989925c755824493ab564a4126a5a344",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 36 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01e8b21cef294df1a87cb5c28ef82ea2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a8e7a39373d42e882b70b52b5433262",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 37 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4beec5f4b2d74fdf92c118f2acb176f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e88a87025ede457eb21b0dc386962e6f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 38 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90309fc1a4064cf49ff7a182ce034f3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb9cd6eff36445048c830d3d94cb87e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 39 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea947fce1d5e48fb8c359ace970b9ce5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a297b6254ba04a608f8d6347bdaa8fa8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 40 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0171a170b4aa405a86985b8671ccb219",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2134914f6b94411aac9c6d4efea37c1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 41 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccaf5218863b4014ae07f02565f5d4e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df09a75f2da54406b6d1fb42d10019d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 42 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d4d4d0893194bd986951ba0e7dace1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a65cba18f094fef8215a784b7b33032",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 43 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be4a505e8274431bba050706df48aec7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cddbb1508fa3404a9b3ae0923f541384",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 44 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "290d7950a6884d97a0a22480525097f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ddc413606b544d6497a5d2d8a26c0abd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 45 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80bb2576c07c461d93d130bf50f31447",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07265c6e3c2244979f17bac91e38fa7b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 46 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09d8a5f325764b7a9c2d0e0609198c82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a558f73744248489c9d1d487eada45d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 47 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3348a93b81894917800ae3981e567b87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ec825f7fd684f75b4c924e188777f82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 48 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9bd94534d4564f99b76be3439e90c8cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "383380cf00474a7c8a309e87a94fb3e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 49 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95d83375e9e74f00bc6b253c14c62c65",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbf1c687237d42a095793e79d0d68897",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 50 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e72c94d14c1146b09d0bb6bb20b888ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42545b91b47f49189081609f32235035",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 51 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f9a844aebb04d799dfc865f0dc7a1db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "256df1d8d0fa41b3973b9d4a1bc3f7e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 52 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b249a07abe541d5acc4b66402142b83",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "203711350abf4ab68cc31f2252013a4a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 53 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b87c2d7d628d43aaa1d01857802c8dc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f3734ae79a19455ea2a8f64f68fe8010",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 54 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee077ff4771044ce88b165ac3c8e883a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd0cc92b24764af1a83e90edbbc6b51c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 55 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4c7265a507340cf8cee16fc7e29a187",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7218f2325c494354a509f199d1d1b78f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 56 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f3b9fb14cf94e36baaf991d34e14c7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "582b0c2768334eecaa41dc1177d6986d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 57 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22c359d8b4b74fa0a95dfc183e4a5e0a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61d9f65b5f174b8b8e2b8e9a2235afb2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 58 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2204aaa6778c470598f7b2f19f0c5270",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6061b33d87254ee898d7b2cbb957a1d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== EPOCH 59 ===========\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d76398811d94494aea3160fd6f9f756",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae9c1a1956294f6cb7870635687c4150",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "73d1f6cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-bddc8ef30e891a65\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-bddc8ef30e891a65\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir tensorboard/ # Change for your log location"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24a2cefe",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "20269d37",
   "metadata": {},
   "outputs": [],
   "source": [
    "#trainer.resume_model(49)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3a4afa0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99b71a64a644405d8001aae5bee5c771",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saved prediction las files\n"
     ]
    }
   ],
   "source": [
    "trainer.test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5df27a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rock_detection_3d.utils.las_reader import Read_Las_from_Path, Read_Las_from_Json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c665e2fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n"
     ]
    }
   ],
   "source": [
    "pred_las_reader = Read_Las_from_Path('data/rocklas/prediction_3d')\n",
    "print(len(pred_las_reader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f6130eaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n"
     ]
    }
   ],
   "source": [
    "las_reader = Read_Las_from_Json('data/rocklas/raw/test_split.json')\n",
    "print(len(las_reader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "fa13da1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a4309a1691f4dd0a06a94036ddc72d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ViewInteractiveWidget(height=412, layout=Layout(height='auto', width='100%'), width=1024)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx = 6\n",
    "\n",
    "p = pv.Plotter(notebook=True,shape=(1, 2),window_size=[1024,412])\n",
    "\n",
    "p.subplot(0, 0)\n",
    "pos, color, y = pred_las_reader.get_normalized(idx)\n",
    "point_cloud = pv.PolyData(pos)\n",
    "point_cloud['y'] = y\n",
    "p.add_points(point_cloud,  show_scalar_bar=False, point_size=4)\n",
    "p.camera_position = [-1,5, -10]\n",
    "               \n",
    "p.subplot(0, 1)\n",
    "pos, color, y = las_reader.get_normalized(idx)\n",
    "point_cloud = pv.PolyData(pos)\n",
    "point_cloud['y'] = y\n",
    "p.add_points(point_cloud,  show_scalar_bar=False, point_size=4)\n",
    "p.camera_position = [-1,5, -10]          \n",
    "               \n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fb7ad96",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}