inference_cylinder_asym_jupyter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "ba7e7c9e-33e7-4320-b0ae-e0977dcddae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import argparse\n",
    "\n",
    "import os\n",
    "import sys\n",
    "sys.path.append('.')\n",
    "\n",
    "import yaml\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_scatter\n",
    "import spconv.pytorch as spconv\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from utils.load_save_util import load_checkpoint\n",
    "\n",
    "from config.config import load_config_data\n",
    "from utils.lovasz_losses import lovasz_softmax\n",
    "\n",
    "from dataloader.dataset_semantickitti import get_model_class, collate_fn_BEV\n",
    "from dataloader.pc_dataset import get_pc_model_class\n",
    "from dataloader.pc_dataset import get_SemKITTI_label_name\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a33bccb2-456a-42e8-9542-ba2e143987db",
   "metadata": {},
   "source": [
    "## Cylinder Feature Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "0f0df192-7745-4394-bd49-f3ef8a887583",
   "metadata": {},
   "outputs": [],
   "source": [
    "class cylinder_fea(nn.Module):\n",
    "\n",
    "    def __init__(self, grid_size, fea_dim=3,\n",
    "                 out_pt_fea_dim=64, max_pt_per_encode=64, fea_compre=None):\n",
    "        super(cylinder_fea, self).__init__()\n",
    "\n",
    "        self.PPmodel = nn.Sequential(\n",
    "            nn.BatchNorm1d(fea_dim),\n",
    "\n",
    "            nn.Linear(fea_dim, 64),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Linear(64, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Linear(128, 256),\n",
    "            nn.BatchNorm1d(256),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Linear(256, out_pt_fea_dim)\n",
    "        )\n",
    "\n",
    "        self.max_pt = max_pt_per_encode\n",
    "        self.fea_compre = fea_compre\n",
    "        self.grid_size = grid_size\n",
    "        kernel_size = 3\n",
    "        self.local_pool_op = torch.nn.MaxPool2d(kernel_size, stride=1,\n",
    "                                                padding=(kernel_size - 1) // 2,\n",
    "                                                dilation=1)\n",
    "        self.pool_dim = out_pt_fea_dim\n",
    "\n",
    "        # point feature compression\n",
    "        if self.fea_compre is not None:\n",
    "            self.fea_compression = nn.Sequential(\n",
    "                nn.Linear(self.pool_dim, self.fea_compre),\n",
    "                nn.ReLU())\n",
    "            self.pt_fea_dim = self.fea_compre\n",
    "        else:\n",
    "            self.pt_fea_dim = self.pool_dim\n",
    "\n",
    "    def forward(self, pt_fea, xy_ind):\n",
    "        cur_dev = pt_fea[0].get_device()\n",
    "\n",
    "        # concate everything\n",
    "        cat_pt_ind = []\n",
    "        for i_batch in range(len(xy_ind)):\n",
    "            cat_pt_ind.append(F.pad(xy_ind[i_batch], (1, 0), 'constant', value=i_batch))\n",
    "\n",
    "        cat_pt_fea = torch.cat(pt_fea, dim=0)\n",
    "        cat_pt_ind = torch.cat(cat_pt_ind, dim=0)\n",
    "        pt_num = cat_pt_ind.shape[0]\n",
    "\n",
    "        # shuffle the data\n",
    "        shuffled_ind = torch.randperm(pt_num, device=cur_dev)\n",
    "        cat_pt_fea = cat_pt_fea[shuffled_ind, :]\n",
    "        cat_pt_ind = cat_pt_ind[shuffled_ind, :]\n",
    "\n",
    "        # unique xy grid index\n",
    "        unq, unq_inv, unq_cnt = torch.unique(cat_pt_ind, return_inverse=True, return_counts=True, dim=0)\n",
    "        unq = unq.type(torch.int64)\n",
    "\n",
    "        # process feature\n",
    "        processed_cat_pt_fea = self.PPmodel(cat_pt_fea)\n",
    "        pooled_data = torch_scatter.scatter_max(processed_cat_pt_fea, unq_inv, dim=0)[0]\n",
    "\n",
    "        if self.fea_compre:\n",
    "            processed_pooled_data = self.fea_compression(pooled_data)\n",
    "        else:\n",
    "            processed_pooled_data = pooled_data\n",
    "\n",
    "        return unq, processed_pooled_data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13bd1c96-2426-4e94-86b5-0f6b8a5ea419",
   "metadata": {},
   "source": [
    "## Cylinder Sparse Convolution 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "cbcd2bb5-c2ea-4da6-8665-78b3630b5a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "REGISTERED_MODELS_CLASSES = {}\n",
    "\n",
    "def register_model(cls, name=None):\n",
    "    global REGISTERED_MODELS_CLASSES\n",
    "    if name is None:\n",
    "        name = cls.__name__\n",
    "    assert name not in REGISTERED_MODELS_CLASSES, f\"exist class: {REGISTERED_MODELS_CLASSES}\"\n",
    "    REGISTERED_MODELS_CLASSES[name] = cls\n",
    "    return cls\n",
    "\n",
    "\n",
    "def get_model_class_c3d(name):\n",
    "    global REGISTERED_MODELS_CLASSES\n",
    "    assert name in REGISTERED_MODELS_CLASSES, f\"available class: {REGISTERED_MODELS_CLASSES}\"\n",
    "    return REGISTERED_MODELS_CLASSES[name]\n",
    "\n",
    "\n",
    "@register_model\n",
    "class cylinder_asym(nn.Module):\n",
    "    def __init__(self,\n",
    "                 cylin_model,\n",
    "                 segmentator_spconv,\n",
    "                 sparse_shape,\n",
    "                 ):\n",
    "        super().__init__()\n",
    "        self.name = \"cylinder_asym\"\n",
    "\n",
    "        self.cylinder_3d_generator = cylin_model\n",
    "\n",
    "        self.cylinder_3d_spconv_seg = segmentator_spconv\n",
    "\n",
    "        self.sparse_shape = sparse_shape\n",
    "\n",
    "    def forward(self, train_pt_fea_ten, train_vox_ten, batch_size):\n",
    "        coords, features_3d = self.cylinder_3d_generator(train_pt_fea_ten, train_vox_ten)\n",
    "\n",
    "        spatial_features = self.cylinder_3d_spconv_seg(features_3d, coords, batch_size)\n",
    "\n",
    "        return spatial_features\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a1cd19-6cda-4d07-bc0e-8f3318378fe1",
   "metadata": {},
   "source": [
    "## Segmentator 3D Asymmetric Sparse Convolution 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "524dae7b-58c8-43bc-a47b-64c64e804cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv3x3(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                             padding=1, bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv1x3(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=(1, 3, 3), stride=stride,\n",
    "                             padding=(0, 1, 1), bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv1x1x3(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=(1, 1, 3), stride=stride,\n",
    "                             padding=(0, 0, 1), bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv1x3x1(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=(1, 3, 1), stride=stride,\n",
    "                             padding=(0, 1, 0), bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv3x1x1(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=(3, 1, 1), stride=stride,\n",
    "                             padding=(1, 0, 0), bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv3x1(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=(3, 1, 3), stride=stride,\n",
    "                             padding=(1, 0, 1), bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1, indice_key=None):\n",
    "    return spconv.SubMConv3d(in_planes, out_planes, kernel_size=1, stride=stride,\n",
    "                             padding=1, bias=False, indice_key=indice_key)\n",
    "\n",
    "\n",
    "class ResContextBlock(nn.Module):\n",
    "    def __init__(self, in_filters, out_filters, kernel_size=(3, 3, 3), stride=1, indice_key=None):\n",
    "        super(ResContextBlock, self).__init__()\n",
    "        self.conv1 = conv1x3(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.bn0 = nn.BatchNorm1d(out_filters)\n",
    "        self.act1 = nn.LeakyReLU()\n",
    "          \n",
    "        #elf.conv1_2 = conv3x1(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.conv1_2 = conv1x3(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "\n",
    "        self.bn0_2 = nn.BatchNorm1d(out_filters)\n",
    "        self.act1_2 = nn.LeakyReLU()\n",
    "\n",
    "        self.conv2 = conv3x1(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act2 = nn.LeakyReLU()\n",
    "        self.bn1 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        #self.conv3 = conv1x3(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.conv3 = conv3x1(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act3 = nn.LeakyReLU()\n",
    "        self.bn2 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        self.weight_initialization()\n",
    "\n",
    "    def weight_initialization(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.BatchNorm1d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        shortcut = self.conv1(x)\n",
    "        shortcut = shortcut.replace_feature(self.act1(shortcut.features))\n",
    "        shortcut = shortcut.replace_feature(self.bn0(shortcut.features))\n",
    "\n",
    "        shortcut = self.conv1_2(shortcut)\n",
    "        shortcut = shortcut.replace_feature(self.act1_2(shortcut.features))\n",
    "        shortcut = shortcut.replace_feature(self.bn0_2(shortcut.features))\n",
    "\n",
    "        resA = self.conv2(x)\n",
    "        resA = resA.replace_feature(self.act2(resA.features))\n",
    "        reaA = resA.replace_feature(self.bn1(resA.features))\n",
    "\n",
    "        resA = self.conv3(resA)\n",
    "        resA = resA.replace_feature(self.act3(resA.features))\n",
    "        resA = resA.replace_feature(self.bn2(resA.features))\n",
    "        resA = resA.replace_feature(resA.features + shortcut.features)\n",
    "\n",
    "        return resA\n",
    "\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, in_filters, out_filters, dropout_rate, kernel_size=(3, 3, 3), stride=1,\n",
    "                 pooling=True, drop_out=True, height_pooling=False, indice_key=None):\n",
    "        super(ResBlock, self).__init__()\n",
    "        self.pooling = pooling\n",
    "        self.drop_out = drop_out\n",
    "\n",
    "        self.conv1 = conv3x1(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act1 = nn.LeakyReLU()\n",
    "        self.bn0 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        #self.conv1_2 = conv1x3(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.conv1_2 = conv3x1(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act1_2 = nn.LeakyReLU()\n",
    "        self.bn0_2 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        self.conv2 = conv1x3(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act2 = nn.LeakyReLU()\n",
    "        self.bn1 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        #self.conv3 = conv3x1(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.conv3 = conv1x3(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.act3 = nn.LeakyReLU()\n",
    "        self.bn2 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        if pooling:\n",
    "            if height_pooling:\n",
    "                self.pool = spconv.SparseConv3d(out_filters, out_filters, kernel_size=3, stride=2,\n",
    "                                                padding=1, indice_key=indice_key, bias=False)\n",
    "            else:\n",
    "                self.pool = spconv.SparseConv3d(out_filters, out_filters, kernel_size=3, stride=(2, 2, 1),\n",
    "                                                padding=1, indice_key=indice_key, bias=False)\n",
    "        self.weight_initialization()\n",
    "\n",
    "    def weight_initialization(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.BatchNorm1d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        shortcut = self.conv1(x)\n",
    "        shortcut = shortcut.replace_feature(self.act1(shortcut.features))\n",
    "        shortcut = shortcut.replace_feature(self.bn0(shortcut.features))\n",
    "\n",
    "        shortcut = self.conv1_2(shortcut)\n",
    "        shortcut = shortcut.replace_feature(self.act1_2(shortcut.features))\n",
    "        shortcut = shortcut.replace_feature(self.bn0_2(shortcut.features))\n",
    "\n",
    "        resA = self.conv2(x)\n",
    "        resA = resA.replace_feature(self.act2(resA.features))\n",
    "        resA = resA.replace_feature(self.bn1(resA.features))\n",
    "\n",
    "        resA = self.conv3(resA)\n",
    "        resA = resA.replace_feature(self.act3(resA.features))\n",
    "        resA = resA.replace_feature(self.bn2(resA.features))\n",
    "\n",
    "        resA = resA.replace_feature(resA.features + shortcut.features)\n",
    "\n",
    "        if self.pooling:\n",
    "            resB = self.pool(resA)\n",
    "            return resB, resA\n",
    "        else:\n",
    "            return resA\n",
    "\n",
    "\n",
    "class UpBlock(nn.Module):\n",
    "    def __init__(self, in_filters, out_filters, kernel_size=(3, 3, 3), indice_key=None, up_key=None):\n",
    "        super(UpBlock, self).__init__()\n",
    "        # self.drop_out = drop_out\n",
    "        self.trans_dilao = conv3x3(in_filters, out_filters, indice_key=indice_key + \"new_up\")\n",
    "        self.trans_act = nn.LeakyReLU()\n",
    "        self.trans_bn = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        self.conv1 = conv1x3(out_filters, out_filters, indice_key=indice_key)\n",
    "        self.act1 = nn.LeakyReLU()\n",
    "        self.bn1 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        #self.conv3 = conv3x1(out_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.conv2 = conv1x3(out_filters, out_filters, indice_key=indice_key)\n",
    "        self.act2 = nn.LeakyReLU()\n",
    "        self.bn2 = nn.BatchNorm1d(out_filters)\n",
    "\n",
    "        #self.conv3 = conv3x3(out_filters, out_filters, indice_key=indice_key)\n",
    "        self.conv3 = conv1x3(out_filters, out_filters, indice_key=indice_key)\n",
    "        self.act3 = nn.LeakyReLU()\n",
    "        self.bn3 = nn.BatchNorm1d(out_filters)\n",
    "        # self.dropout3 = nn.Dropout3d(p=dropout_rate)\n",
    "\n",
    "        self.up_subm = spconv.SparseInverseConv3d(out_filters, out_filters, kernel_size=3, indice_key=up_key,\n",
    "                                                  bias=False)\n",
    "\n",
    "        self.weight_initialization()\n",
    "\n",
    "    def weight_initialization(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.BatchNorm1d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def forward(self, x, skip):\n",
    "        upA = self.trans_dilao(x)\n",
    "        upA = upA.replace_feature(self.trans_act(upA.features))\n",
    "        upA = upA.replace_feature(self.trans_bn(upA.features))\n",
    "\n",
    "        ## upsample\n",
    "        upA = self.up_subm(upA)\n",
    "\n",
    "        upA = upA.replace_feature(upA.features + skip.features)\n",
    "\n",
    "        upE = self.conv1(upA)\n",
    "        upE = upE.replace_feature(self.act1(upE.features))\n",
    "        upE = upE.replace_feature(self.bn1(upE.features))\n",
    "\n",
    "        upE = self.conv2(upE)\n",
    "        upE = upE.replace_feature(self.act2(upE.features))\n",
    "        upE = upE.replace_feature(self.bn2(upE.features))\n",
    "\n",
    "        upE = self.conv3(upE)\n",
    "        upE = upE.replace_feature(self.act3(upE.features))\n",
    "        upE = upE.replace_feature(self.bn3(upE.features))\n",
    "\n",
    "        return upE\n",
    "\n",
    "\n",
    "class ReconBlock(nn.Module):\n",
    "    def __init__(self, in_filters, out_filters, kernel_size=(3, 3, 3), stride=1, indice_key=None):\n",
    "        super(ReconBlock, self).__init__()\n",
    "        self.conv1 = conv3x1x1(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.bn0 = nn.BatchNorm1d(out_filters)\n",
    "        self.act1 = nn.Sigmoid()\n",
    "\n",
    "        self.conv1_2 = conv1x3x1(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.bn0_2 = nn.BatchNorm1d(out_filters)\n",
    "        self.act1_2 = nn.Sigmoid()\n",
    "\n",
    "        self.conv1_3 = conv1x1x3(in_filters, out_filters, indice_key=indice_key + \"bef\")\n",
    "        self.bn0_3 = nn.BatchNorm1d(out_filters)\n",
    "        self.act1_3 = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        shortcut = self.conv1(x)\n",
    "        shortcut = shortcut.replace_feature(self.bn0(shortcut.features))\n",
    "        shortcut = shortcut.replace_feature(self.act1(shortcut.features))\n",
    "\n",
    "        shortcut2 = self.conv1_2(x)\n",
    "        shortcut2 = shortcut2.replace_feature(self.bn0_2(shortcut2.features))\n",
    "        shortcut2 = shortcut2.replace_feature(self.act1_2(shortcut2.features))\n",
    "\n",
    "        shortcut3 = self.conv1_3(x)\n",
    "        shortcut3 = shortcut.replace_feature(self.bn0_3(shortcut3.features))\n",
    "        shortcut3 = shortcut3.replace_feature(self.act1_3(shortcut3.features))\n",
    "        shortcut = shortcut.replace_feature(shortcut.features + shortcut2.features + shortcut3.features)\n",
    "\n",
    "        shortcut = shortcut.replace_feature(shortcut.features * x.features)\n",
    "\n",
    "        return shortcut\n",
    "\n",
    "\n",
    "class Asymm_3d_spconv(nn.Module):\n",
    "    def __init__(self,\n",
    "                 output_shape,\n",
    "                 use_norm=True,\n",
    "                 num_input_features=128,\n",
    "                 nclasses=20, n_height=32, strict=False, init_size=16):\n",
    "        super(Asymm_3d_spconv, self).__init__()\n",
    "        self.nclasses = nclasses\n",
    "        self.nheight = n_height\n",
    "        self.strict = False\n",
    "\n",
    "        sparse_shape = np.array(output_shape)\n",
    "        # sparse_shape[0] = 11\n",
    "        print(sparse_shape)\n",
    "        self.sparse_shape = sparse_shape\n",
    "\n",
    "        self.downCntx = ResContextBlock(num_input_features, init_size, indice_key=\"pre\")\n",
    "        self.resBlock2 = ResBlock(init_size, 2 * init_size, 0.2, height_pooling=True, indice_key=\"down2\")\n",
    "        self.resBlock3 = ResBlock(2 * init_size, 4 * init_size, 0.2, height_pooling=True, indice_key=\"down3\")\n",
    "        self.resBlock4 = ResBlock(4 * init_size, 8 * init_size, 0.2, pooling=True, height_pooling=False,\n",
    "                                  indice_key=\"down4\")\n",
    "        self.resBlock5 = ResBlock(8 * init_size, 16 * init_size, 0.2, pooling=True, height_pooling=False,\n",
    "                                  indice_key=\"down5\")\n",
    "\n",
    "        self.upBlock0 = UpBlock(16 * init_size, 16 * init_size, indice_key=\"up0\", up_key=\"down5\")\n",
    "        self.upBlock1 = UpBlock(16 * init_size, 8 * init_size, indice_key=\"up1\", up_key=\"down4\")\n",
    "        self.upBlock2 = UpBlock(8 * init_size, 4 * init_size, indice_key=\"up2\", up_key=\"down3\")\n",
    "        self.upBlock3 = UpBlock(4 * init_size, 2 * init_size, indice_key=\"up3\", up_key=\"down2\")\n",
    "\n",
    "        self.ReconNet = ReconBlock(2 * init_size, 2 * init_size, indice_key=\"recon\")\n",
    "\n",
    "        self.logits = spconv.SubMConv3d(4 * init_size, nclasses, indice_key=\"logit\", kernel_size=3, stride=1, padding=1,\n",
    "                                        bias=True)\n",
    "\n",
    "    def forward(self, voxel_features, coors, batch_size):\n",
    "        # x = x.contiguous()\n",
    "        coors = coors.int()\n",
    "        # import pdb\n",
    "        # pdb.set_trace()\n",
    "        ret = spconv.SparseConvTensor(voxel_features, coors, self.sparse_shape,\n",
    "                                      batch_size)\n",
    "        ret = self.downCntx(ret)\n",
    "        down1c, down1b = self.resBlock2(ret)\n",
    "        down2c, down2b = self.resBlock3(down1c)\n",
    "        down3c, down3b = self.resBlock4(down2c)\n",
    "        down4c, down4b = self.resBlock5(down3c)\n",
    "\n",
    "        up4e = self.upBlock0(down4c, down4b)\n",
    "        up3e = self.upBlock1(up4e, down3b)\n",
    "        up2e = self.upBlock2(up3e, down2b)\n",
    "        up1e = self.upBlock3(up2e, down1b)\n",
    "\n",
    "        up0e = self.ReconNet(up1e)\n",
    "\n",
    "        up0e = up0e.replace_feature(torch.cat((up0e.features, up1e.features), 1))\n",
    "\n",
    "        logits = self.logits(up0e)\n",
    "        y = logits.dense()\n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da70259b-b4c1-45ea-ab49-bafa8dd531c0",
   "metadata": {},
   "source": [
    "## Loss Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "eda89004-b68b-4680-be3b-2cf4b67880cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class loss_builder:\n",
    "    def build(wce=True, lovasz=True, num_class=20, ignore_label=0):\n",
    "\n",
    "        loss_funs = torch.nn.CrossEntropyLoss(ignore_index=ignore_label)\n",
    "\n",
    "        if wce and lovasz:\n",
    "            return loss_funs, lovasz_softmax\n",
    "        elif wce and not lovasz:\n",
    "            return wce\n",
    "        elif not wce and lovasz:\n",
    "            return lovasz_softmax\n",
    "        else:\n",
    "            raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f8f4b0f-2e2c-4da3-b003-dd5c51d7d26c",
   "metadata": {},
   "source": [
    "## Data Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "8495a140-34cf-4e8d-bee7-f3fc3764d49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class data_builder:\n",
    "    def build(dataset_config,\n",
    "              train_dataloader_config,\n",
    "              val_dataloader_config,\n",
    "              grid_size=[480, 360, 32]):\n",
    "        data_path = train_dataloader_config[\"data_path\"]\n",
    "        train_imageset = train_dataloader_config[\"imageset\"]\n",
    "        val_imageset = val_dataloader_config[\"imageset\"]\n",
    "        train_ref = train_dataloader_config[\"return_ref\"]\n",
    "        val_ref = val_dataloader_config[\"return_ref\"]\n",
    "\n",
    "        label_mapping = dataset_config[\"label_mapping\"]\n",
    "\n",
    "        SemKITTI = get_pc_model_class(dataset_config['pc_dataset_type'])\n",
    "\n",
    "        nusc=None\n",
    "        if \"nusc\" in dataset_config['pc_dataset_type']:\n",
    "            from nuscenes import NuScenes\n",
    "            nusc = NuScenes(version='v1.0-trainval', dataroot=data_path, verbose=True)\n",
    "\n",
    "        train_pt_dataset = SemKITTI(data_path, imageset=train_imageset,\n",
    "                                    return_ref=train_ref, label_mapping=label_mapping, nusc=nusc)\n",
    "        val_pt_dataset = SemKITTI(data_path, imageset=val_imageset,\n",
    "                                  return_ref=val_ref, label_mapping=label_mapping, nusc=nusc)\n",
    "\n",
    "        train_dataset = get_model_class_dataset(dataset_config['dataset_type'])(\n",
    "            train_pt_dataset,\n",
    "            grid_size=grid_size,\n",
    "            flip_aug=True,\n",
    "            fixed_volume_space=dataset_config['fixed_volume_space'],\n",
    "            max_volume_space=dataset_config['max_volume_space'],\n",
    "            min_volume_space=dataset_config['min_volume_space'],\n",
    "            ignore_label=dataset_config[\"ignore_label\"],\n",
    "            rotate_aug=True,\n",
    "            scale_aug=True,\n",
    "            transform_aug=True\n",
    "        )\n",
    "\n",
    "        val_dataset = get_model_class_dataset(dataset_config['dataset_type'])(\n",
    "            val_pt_dataset,\n",
    "            grid_size=grid_size,\n",
    "            fixed_volume_space=dataset_config['fixed_volume_space'],\n",
    "            max_volume_space=dataset_config['max_volume_space'],\n",
    "            min_volume_space=dataset_config['min_volume_space'],\n",
    "            ignore_label=dataset_config[\"ignore_label\"],\n",
    "        )\n",
    "\n",
    "        train_dataset_loader = torch.utils.data.DataLoader(dataset=train_dataset,\n",
    "                                                           batch_size=train_dataloader_config[\"batch_size\"],\n",
    "                                                           collate_fn=collate_fn_BEV,\n",
    "                                                           shuffle=train_dataloader_config[\"shuffle\"],\n",
    "                                                           num_workers=train_dataloader_config[\"num_workers\"])\n",
    "        val_dataset_loader = torch.utils.data.DataLoader(dataset=val_dataset,\n",
    "                                                         batch_size=val_dataloader_config[\"batch_size\"],\n",
    "                                                         collate_fn=collate_fn_BEV,\n",
    "                                                         shuffle=val_dataloader_config[\"shuffle\"],\n",
    "                                                         num_workers=val_dataloader_config[\"num_workers\"])\n",
    "\n",
    "        return train_dataset_loader, val_dataset_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f46b44c7-3d2b-40ca-8429-4e58fdb84b59",
   "metadata": {},
   "source": [
    "## Model Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "caf25a09-d6a1-4bd5-b22a-45b70563972a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class model_builder:\n",
    "    def build(model_config):\n",
    "        output_shape = model_config['output_shape']\n",
    "        num_class = model_config['num_class']\n",
    "        num_input_features = model_config['num_input_features']\n",
    "        use_norm = model_config['use_norm']\n",
    "        init_size = model_config['init_size']\n",
    "        fea_dim = model_config['fea_dim']\n",
    "        out_fea_dim = model_config['out_fea_dim']\n",
    "\n",
    "        cylinder_3d_spconv_seg = Asymm_3d_spconv(\n",
    "            output_shape=output_shape,\n",
    "            use_norm=use_norm,\n",
    "            num_input_features=num_input_features,\n",
    "            init_size=init_size,\n",
    "            nclasses=num_class)\n",
    "\n",
    "        cy_fea_net = cylinder_fea(grid_size=output_shape,\n",
    "                                  fea_dim=fea_dim,\n",
    "                                  out_pt_fea_dim=out_fea_dim,\n",
    "                                  fea_compre=num_input_features)\n",
    "\n",
    "        model = get_model_class_c3d(model_config[\"model_architecture\"])(\n",
    "            cylin_model=cy_fea_net,\n",
    "            segmentator_spconv=cylinder_3d_spconv_seg,\n",
    "            sparse_shape=output_shape\n",
    "        )\n",
    "\n",
    "        return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de301085-aadf-4fbe-8db9-9f9d900a5c67",
   "metadata": {},
   "source": [
    "## Build Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "4c5e7767-abe7-45b5-b3bf-9066c9d6969b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(dataset_config,\n",
    "                  data_dir,\n",
    "                  grid_size=[480, 360, 32],\n",
    "                  demo_label_dir=None):\n",
    "\n",
    "    if demo_label_dir == '':\n",
    "        imageset = \"demo\"\n",
    "    else:\n",
    "        imageset = \"val\"\n",
    "    label_mapping = dataset_config[\"label_mapping\"]\n",
    "\n",
    "    SemKITTI_demo = get_pc_model_class('SemKITTI_demo')\n",
    "\n",
    "    demo_pt_dataset = SemKITTI_demo(data_dir, imageset=imageset,\n",
    "                              return_ref=True, label_mapping=label_mapping, demo_label_path=demo_label_dir)\n",
    "\n",
    "    demo_dataset = get_model_class(dataset_config['dataset_type'])(\n",
    "        demo_pt_dataset,\n",
    "        grid_size=grid_size,\n",
    "        fixed_volume_space=dataset_config['fixed_volume_space'],\n",
    "        max_volume_space=dataset_config['max_volume_space'],\n",
    "        min_volume_space=dataset_config['min_volume_space'],\n",
    "        ignore_label=dataset_config[\"ignore_label\"],\n",
    "    )\n",
    "    demo_dataset_loader = torch.utils.data.DataLoader(dataset=demo_dataset,\n",
    "                                                     batch_size=1,\n",
    "                                                     collate_fn=collate_fn_BEV,\n",
    "                                                     shuffle=False,\n",
    "                                                     num_workers=4)\n",
    "\n",
    "    return demo_dataset_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0fc266f-444b-44ee-ac53-06b532bbea5b",
   "metadata": {},
   "source": [
    "## Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "b2ffd65d-0288-46d6-a9c1-285d8e0e0f89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_hist(pred, label, n):\n",
    "    k = (label >= 0) & (label < n)\n",
    "    bin_count = np.bincount(\n",
    "        n * label[k].astype(int) + pred[k], minlength=n ** 2)\n",
    "    return bin_count[:n ** 2].reshape(n, n)\n",
    "\n",
    "\n",
    "def per_class_iu(hist):\n",
    "    return np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))\n",
    "\n",
    "\n",
    "def fast_hist_crop(output, target, unique_label):\n",
    "    hist = fast_hist(output.flatten(), target.flatten(), np.max(unique_label) + 2)\n",
    "    hist = hist[unique_label + 1, :]\n",
    "    hist = hist[:, unique_label + 1]\n",
    "    return hist\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a70e592b-11a6-491e-930f-c86c64520e8c",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "d086ba0b-972c-44af-9c67-ce0b89504c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference(config_path, demo_folder, demo_label_folder, save_folder):\n",
    "    pytorch_device = torch.device('cuda:0')\n",
    "\n",
    "    configs = load_config_data(config_path)\n",
    "    dataset_config = configs['dataset_params']\n",
    "    \n",
    "    data_dir = demo_folder\n",
    "    demo_label_dir = demo_label_folder\n",
    "    save_dir = save_folder + \"/\"\n",
    "\n",
    "    demo_batch_size = 1\n",
    "    model_config = configs['model_params']\n",
    "    train_hypers = configs['train_params']\n",
    "\n",
    "    grid_size = model_config['output_shape']\n",
    "    num_class = model_config['num_class']\n",
    "    ignore_label = dataset_config['ignore_label']\n",
    "    model_load_path = train_hypers['model_load_path']\n",
    "\n",
    "    SemKITTI_label_name = get_SemKITTI_label_name(dataset_config[\"label_mapping\"])\n",
    "    unique_label = np.asarray(sorted(list(SemKITTI_label_name.keys())))[1:] - 1\n",
    "    unique_label_str = [SemKITTI_label_name[x] for x in unique_label + 1]\n",
    "    my_model = model_builder.build(model_config)\n",
    "    if os.path.exists(model_load_path):\n",
    "        #my_model = load_checkpoint(model_load_path, my_model)\n",
    "        # the original pretrained weights are not loaded properly\n",
    "        # i had to change the load function and only a model can be loaded which was trained with this net\n",
    "        my_model.load_state_dict(torch.load(model_load_path))\n",
    "    else:\n",
    "        print(\"------------------------------------!\")\n",
    "        print(\"Model Checkpoint NOT loaded!\")\n",
    "        print(\"------------------------------------!\")\n",
    "\n",
    "    my_model.to(pytorch_device)\n",
    "    optimizer = optim.Adam(my_model.parameters(), lr=train_hypers[\"learning_rate\"])\n",
    "\n",
    "    loss_func, lovasz_softmax = loss_builder.build(wce=True, lovasz=True,\n",
    "                                                   num_class=num_class, ignore_label=ignore_label)\n",
    "\n",
    "    demo_dataset_loader = build_dataset(dataset_config, data_dir, grid_size=grid_size, demo_label_dir=demo_label_dir)\n",
    "    with open(dataset_config[\"label_mapping\"], 'r') as stream:\n",
    "        semkittiyaml = yaml.safe_load(stream)\n",
    "    inv_learning_map = semkittiyaml['learning_map_inv']\n",
    "\n",
    "    my_model.eval()\n",
    "    hist_list = []\n",
    "    demo_loss_list = []\n",
    "    with torch.no_grad():\n",
    "        for i_iter_demo, (_, demo_vox_label, demo_grid, demo_pt_labs, demo_pt_fea) in enumerate(\n",
    "                demo_dataset_loader):\n",
    "            demo_pt_fea_ten = [torch.from_numpy(i).type(torch.FloatTensor).to(pytorch_device) for i in\n",
    "                              demo_pt_fea]\n",
    "            demo_grid_ten = [torch.from_numpy(i).to(pytorch_device) for i in demo_grid]\n",
    "            demo_label_tensor = demo_vox_label.type(torch.LongTensor).to(pytorch_device)\n",
    "\n",
    "            predict_labels = my_model(demo_pt_fea_ten, demo_grid_ten, demo_batch_size)\n",
    "            loss = lovasz_softmax(torch.nn.functional.softmax(predict_labels).detach(), demo_label_tensor,\n",
    "                                  ignore=0) + loss_func(predict_labels.detach(), demo_label_tensor)\n",
    "            predict_labels = torch.argmax(predict_labels, dim=1)\n",
    "            predict_labels = predict_labels.cpu().detach().numpy()\n",
    "            for count, i_demo_grid in enumerate(demo_grid):\n",
    "                hist_list.append(fast_hist_crop(predict_labels[\n",
    "                                                    count, demo_grid[count][:, 0], demo_grid[count][:, 1],\n",
    "                                                    demo_grid[count][:, 2]], demo_pt_labs[count],\n",
    "                                                unique_label))\n",
    "                inv_labels = np.vectorize(inv_learning_map.__getitem__)(predict_labels[count, demo_grid[count][:, 0], demo_grid[count][:, 1], demo_grid[count][:, 2]]) \n",
    "                inv_labels = inv_labels.astype('uint32')\n",
    "                outputPath = save_dir + str(i_iter_demo).zfill(6) + '.label'\n",
    "                inv_labels.tofile(outputPath)\n",
    "                sys.stdout.write(\"\\rsave \" + outputPath)\n",
    "                sys.stdout.flush()\n",
    "\n",
    "            demo_loss_list.append(loss.detach().cpu().numpy())\n",
    "\n",
    "    if demo_label_dir != '':\n",
    "        my_model.train()\n",
    "        iou = per_class_iu(sum(hist_list))\n",
    "        print('Validation per class iou: ')\n",
    "        for class_name, class_iou in zip(unique_label_str, iou):\n",
    "            print('%s : %.2f%%' % (class_name, class_iou * 100))\n",
    "        val_miou = np.nanmean(iou) * 100\n",
    "        del demo_vox_label, demo_grid, demo_pt_fea, demo_grid_ten\n",
    "\n",
    "        print('Current val miou is %.3f' %\n",
    "              (val_miou))\n",
    "        print('Current val loss is %.3f' %\n",
    "              (np.mean(demo_loss_list))) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce1d6ff6-8057-4d10-b8de-976d71738bbf",
   "metadata": {},
   "source": [
    "## Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bc88aa-b441-433d-bac3-6cd529c57f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference Settings\n",
    "config_path = \"config/semantickitti_jupyter.yaml\"\n",
    "demo_folder = \"../dataset/semantic-kitti/sequences/00/velodyne/\"\n",
    "save_folder = \"save_folder\"\n",
    "demo_label_folder = \"../dataset/semantic-kitti/sequences/00/labels/\"\n",
    "\n",
    "# Inference\n",
    "inference(config_path, demo_folder, demo_label_folder, save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "274d153e-4c65-497b-b1dc-a9fdabf74b3b",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}