predict_labels_attentivecylinder3d.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7e7c9e-33e7-4320-b0ae-e0977dcddae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import csv \n",
    "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",
    "import numba as nb\n",
    "from tqdm import tqdm\n",
    "import multiprocessing\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Function\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",
    "import MinkowskiEngine as ME\n",
    "from utils.metric_util import per_class_iu, fast_hist_crop\n",
    "from utils.lovasz_losses import lovasz_softmax\n",
    "\n",
    "from dataloader.pc_dataset import get_SemKITTI_label_name\n",
    "from dataloader.dataset_semantickitti import get_model_class as get_model_class_dataset, collate_fn_BEV\n",
    "from dataloader.pc_dataset import get_pc_model_class\n",
    "\n",
    "from config.config import load_config_data\n",
    "\n",
    "from utils.load_save_util import load_checkpoint\n",
    "\n",
    "from models_.modules.common import ConvType, NormType, get_norm, conv, get_nonlinearity_fn\n",
    "from models_.modules.resnet_block import *\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "#print(torch.__version__)\n",
    "pytorch_device = torch.device('cuda:0')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13bd1c96-2426-4e94-86b5-0f6b8a5ea419",
   "metadata": {},
   "source": [
    "## Cylinder Sparse Convolution 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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             # Network Architecture - First Part\n",
    "        self.cylinder_3d_spconv_seg = segmentator_spconv     # Network Architecture - Second Part\n",
    "\n",
    "        self.sparse_shape = sparse_shape\n",
    "\n",
    "    def forward(self, train_pt_fea_ten, train_vox_ten, batch_size, iter):\n",
    "        coords, features_3d = self.cylinder_3d_generator(train_pt_fea_ten, train_vox_ten)   # Network Architecture - First Part\n",
    "        spatial_features = self.cylinder_3d_spconv_seg(features_3d, coords, batch_size, iter)     # Network Architecture - Second Part\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": null,
   "id": "524dae7b-58c8-43bc-a47b-64c64e804cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################################################################################\n",
    "# ASYMMETRIC CONVOLUTIONS                                                                #\n",
    "##########################################################################################\n",
    "\n",
    "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",
    "##########################################################################################\n",
    "# ResContextBlock - Used before going into the first DownSample Block.                   #\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",
    "        #self.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",
    "##########################################################################################\n",
    "# ResBlock - Used to create DownSample Blocks. There are originally 4 DownSample Blocks. #\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",
    "##########################################################################################\n",
    "# UpSample Block. There are originally 4 UpSample Blocks                                 #\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",
    "##########################################################################################\n",
    "# DDCM - Figure 4 (right hand side) and architecture Figure: after last U block          #\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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff64a62f",
   "metadata": {},
   "source": [
    "## Transformer Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b3a9018",
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################################################################################\n",
    "# Transformer Block from CodedVTR                                                        #\n",
    "##########################################################################################\n",
    "\n",
    "class CodedVTRBlock(nn.Module): # ddp could not contain unused parameter, so donnot inherit from TRBlock\n",
    "    expansion=1\n",
    "    def __init__(self,\n",
    "               inplanes,\n",
    "               planes,\n",
    "               stride=1,\n",
    "               dilation=1,\n",
    "               downsample=None,\n",
    "               conv_type=ConvType.HYPERCUBE,\n",
    "               nonlinearity_type='ReLU',\n",
    "               bn_momentum=0.1,\n",
    "               D=3):\n",
    "\n",
    "        super(CodedVTRBlock, self).__init__()\n",
    "\n",
    "        self.inplanes = inplanes\n",
    "        self.planes = planes\n",
    "        '''\n",
    "        The Codebook-based Attention: for original feature [1, dim], generate the attnmap [K,h];\n",
    "        then do dotproduct with codebook [D,M,K,dim], get choice [D,M], use it to aggregate the codebook;\n",
    "        apply on value [1, dim] to generate final feature\n",
    "        ------------------\n",
    "        inplanes/outplanes: the feature dim\n",
    "        expansion: the width expansion\n",
    "        qk_type:\n",
    "            - conv\n",
    "            - dotproduct(pairwise)\n",
    "        conv_v: use conv or linear for gen value\n",
    "        vec_dim: the attn_map feature dim\n",
    "        H: head num\n",
    "        D,M - codebook size\n",
    "        K - neighbor-size\n",
    "\n",
    "        The Geometry-based Attention\n",
    "        ------------------\n",
    "        custom-kernel: use CROSS-like / different dilations of neighbor\n",
    "        geo-shape: whether apply geo-shape for codebook elements\n",
    "        temp - the softmax temperature\n",
    "        ------------------\n",
    "        '''\n",
    "\n",
    "        self.expansion = 2\n",
    "        self.qk_type = 'conv' # ['conv','pairwise']\n",
    "        self.conv_v = True\n",
    "        self.top_k_choice = False\n",
    "        self.temp_ = 2.e0 # the initial temp\n",
    "\n",
    "        # === some additonal tricks ===\n",
    "        self.skip_choice = False # only_used in debug mode, notice that this mode contains unused params, so could not support ddp for now\n",
    "        self.geo_shape = False # used in v5 - no improvement, slightly worse than v3\n",
    "        self.sparse_pattern_reg = False # used in v5 - no improvement, slightly worse than v3\n",
    "\n",
    "        if self.inplanes != self.planes:\n",
    "            self.linear_top = MinkoskiConvBNReLU(inplanes, planes, kernel_size=1)\n",
    "            self.downsample = ME.MinkowskiConvolution(inplanes, planes, kernel_size=1, dimension=3)\n",
    "\n",
    "        if self.conv_v == True:\n",
    "            self.v = nn.Sequential(\n",
    "                    MinkoskiConvBNReLU(planes, planes, kernel_size=3),\n",
    "                    MinkoskiConvBNReLU(planes, planes*self.expansion, kernel_size=1),\n",
    "                    )\n",
    "        else:\n",
    "            self.v = MinkoskiConvBNReLU(planes, planes*self.expansion, kernel_size=1)\n",
    "\n",
    "        self.codebook = nn.ModuleList([])\n",
    "        self.D = 3\n",
    "        self.M = 8\n",
    "        self.CUSTOM_KERNEL = True\n",
    "        if self.CUSTOM_KERNEL:  # distinct geometric shape for codebook elements\n",
    "            kgargs0 = {\n",
    "                \"kernel_size\": 3,\n",
    "                \"stride\": 1,\n",
    "                \"dilation\": 2,\n",
    "                # \"region_type\":ME.RegionType.HYPER_CROSS,\n",
    "                \"region_type\":ME.RegionType.HYPER_CUBE,\n",
    "                \"dimension\": 3,\n",
    "                }\n",
    "            kgargs1 = {\n",
    "                \"kernel_size\": 3,\n",
    "                \"stride\": 1,\n",
    "                \"dilation\": 1,\n",
    "                \"region_type\":ME.RegionType.HYPER_CUBE,\n",
    "                \"dimension\": 3,\n",
    "                }\n",
    "            kgargs2 = {\n",
    "                \"kernel_size\": 3,\n",
    "                \"stride\": 1,\n",
    "                \"dilation\": 3,\n",
    "                \"region_type\":ME.RegionType.HYPER_CUBE,\n",
    "                \"dimension\": 3,\n",
    "                }\n",
    "            self.kgargs = [kgargs0, kgargs1, kgargs2] # len should align with M\n",
    "            kgs = [ME.KernelGenerator(**kg) for kg in self.kgargs]\n",
    "            for i_ in range(self.D):\n",
    "                self.codebook.append(\n",
    "                    nn.Sequential(\n",
    "                        ME.MinkowskiChannelwiseConvolution(planes*self.expansion, kernel_size=3, dimension=3, kernel_generator=kgs[i_]),\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "            if not self.skip_choice:\n",
    "                if self.qk_type == 'conv':\n",
    "                    self.q = nn.ModuleList([])\n",
    "                    for i_ in range(self.D):\n",
    "                        self.q.append(\n",
    "                            nn.Sequential(\n",
    "                                ME.MinkowskiConvolution(planes,self.M, dimension=3, kernel_generator=kgs[i_]),\n",
    "                                )\n",
    "                            )\n",
    "                elif self.qk_type == 'pairwise':\n",
    "                    self.q = MinkoskiConvBNReLU(planes, self.M, kernel_size=1)\n",
    "                    # self.pos_enc = MinkoskiConvBNReLU(3, self.M, kernel_size=1)\n",
    "\n",
    "        else:\n",
    "            kgargs0 = {\n",
    "                \"kernel_size\": 3,\n",
    "                \"stride\": 1,\n",
    "                \"dilation\": 1,\n",
    "                \"region_type\":ME.RegionType.HYPER_CUBE,\n",
    "                \"dimension\": 3,\n",
    "                }\n",
    "            self.kgargs = [kgargs0]*self.D\n",
    "            for i_ in range(self.D):\n",
    "                self.codebook.append(\n",
    "                    nn.Sequential(\n",
    "                        ME.MinkowskiConvolution(planes,self.M, dimension=3, kernel_generator=kgs[i_]),\n",
    "                        )\n",
    "                    )\n",
    "            if not self.skip_choice:\n",
    "                if self.qk_type == 'conv':  # since conv already contains the neighbor info, so no pos_enc\n",
    "                    self.q = nn.Sequential(\n",
    "                        ME.MinkowskiConvolution(planes, self.M, kernel_size=3,dimension=3),\n",
    "                        )\n",
    "                elif self.qk_type == 'pairwise':\n",
    "                    self.q = MinkoskiConvBNReLU(planes, self.M, kernel_size=1)\n",
    "                else:\n",
    "                    raise NotImplementedError\n",
    "\n",
    "        if self.geo_shape:\n",
    "            # 3 masks\n",
    "            # each contains masks at differnt stride\n",
    "            # mask1 = torch.load('./plot/final/sparse_masks.pth')\n",
    "            mask0 = np.array([\n",
    "                    [0,1,3,6,7,13],\n",
    "                    [1,2,9,14,15,17],\n",
    "                    [0,5,6,7,8,10],\n",
    "                    [17,19,20,22,23],\n",
    "                    ])\n",
    "            mask1 = np.array([\n",
    "                    [10,11,12,20,21,22],\n",
    "                    [1,2,3,4,5,6,10,21,20],\n",
    "                    [3,4,5,6,7,8,9,10,11],\n",
    "                    [17,18,19,20,22,23,24],\n",
    "                    ])\n",
    "            mask2 = np.array([\n",
    "                    [0,5,9,13,19,22],\n",
    "                    [1,3,7,8,11,16,20],\n",
    "                    [4,6,11,12,18,24,25],\n",
    "                    [5,6,10,14,19,23],\n",
    "                    ])\n",
    "\n",
    "            self.codebook_masks = [mask0, mask1, mask2]\n",
    "            \n",
    "            \n",
    "            for _ in range(len(self.codebook)):\n",
    "                new_kernel = self.codebook[_][0].kernel\n",
    "                k_, dim_ = new_kernel.shape\n",
    "                if len(self.codebook_masks[_])>0:\n",
    "                    assert self.M % len(self.codebook_masks[_]) == 0\n",
    "                if len(self.codebook_masks[_])>1:\n",
    "                    dim_per_mask = dim_ // len(self.codebook_masks[_])\n",
    "                else:\n",
    "                    dim_per_mask = dim_\n",
    "                for m_ in range(len(self.codebook_masks[_])):\n",
    "                    new_kernel[self.codebook_masks[_][m_],dim_per_mask*m_:dim_per_mask*(m_+1)] = 0\n",
    "                self.codebook[_][0].kernel = nn.Parameter(self.codebook[_][0].kernel)\n",
    "\n",
    "            # codebook_weight = torch.stack([m[0].kernel for m in self.codebook])\n",
    "\n",
    "        self.out_bn_relu = nn.Sequential(\n",
    "                ME.MinkowskiConvolution(planes*self.expansion, planes, kernel_size=1, dimension=3),\n",
    "                ME.MinkowskiBatchNorm(planes),\n",
    "                ME.MinkowskiReLU(),\n",
    "                )\n",
    "\n",
    "    def expand_dim(self,x):\n",
    "        # x shold be like [N, vec_dim]; [N, vec_dim, M]\n",
    "        # expand em as [N, dim]; [N, dim, M]\n",
    "        assert x.shape[1] == self.M\n",
    "        if len(x.shape) == 2:\n",
    "            N, dim = x.shape\n",
    "            x = x.unsqueeze(2).expand(-1,-1,self.planes*self.expansion//self.M).reshape(-1,self.planes*self.expansion)\n",
    "        elif len(x.shape) == 3:\n",
    "            N, dim, M = x.shape\n",
    "            x = x.unsqueeze(2).expand(-1,-1,self.planes*self.expansion//self.M, -1).reshape(-1,self.planes*self.expansion,M)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def get_sparse_pattern(self, x, choice, type_=1):\n",
    "        # FORMULA 1: get codebook kernel shapes and directly use the sparse-pattern matching \n",
    "        # as the guidance of choice\n",
    "        if type_ == 1:\n",
    "\n",
    "            sparse_patterns= []  # [M]\n",
    "            for m_ in range(self.D):\n",
    "                kgargs = self.kgargs[m_]\n",
    "                if 'dimension' in kgargs.keys():\n",
    "                    del kgargs['dimension']\n",
    "                neis_d = x.coordinate_manager.get_kernel_map(x.coordinate_map_key,\n",
    "                                                                    x.coordinate_map_key,\n",
    "                                                                    **kgargs\n",
    "                                                                    )\n",
    "                N = x.C.shape[0]\n",
    "                # its easy to get how many matched elements of cur-point & kernel\n",
    "                # but the kernel shape is hard to be flexible, like i need to index the lower-right part\n",
    "                if self.geo_shape:\n",
    "                    # only when codebook-prior is given, each point would have different pattern\n",
    "                    sparse_pattern_ = torch.zeros([N, self.M], device=x.device)\n",
    "                else:\n",
    "                    sparse_pattern_ = torch.zeros([N, 1], device=x.device)\n",
    "\n",
    "                if hasattr(self, \"codebook_masks\"):\n",
    "                    # TODO: acquire the stride corresponding \n",
    "                    cur_mask = self.codebook_masks[m_]\n",
    "                else:\n",
    "                    cur_mask = []\n",
    "\n",
    "                cur_k = len(neis_d.keys())\n",
    "                for k_ in range(cur_k):\n",
    "\n",
    "                    if not k_ in neis_d.keys():\n",
    "                            continue\n",
    "\n",
    "                    if len(cur_mask)>0:\n",
    "                        for i_ in range(len(cur_mask)):\n",
    "                            if k_ in cur_mask[i_]:  # for masked k\n",
    "                                continue\n",
    "                            else:\n",
    "                                sparse_pattern_[neis_d[k_][0].long(),i_] +=1\n",
    "                    else:\n",
    "                        sparse_pattern_[neis_d[k_][0].long(),:] +=1\n",
    "\n",
    "                if len(cur_mask)>0:\n",
    "                    for i_ in range(len(cur_mask)):\n",
    "                        sparse_pattern_[:,i_] = sparse_pattern_[:,i_] / (cur_k - len(cur_mask[i_]))\n",
    "                        if cur_k == len(cur_mask[i_]): # assert zero division, empty kernel\n",
    "                            import ipdb; ipdb.set_trace()\n",
    "                else:\n",
    "                    sparse_pattern_ = sparse_pattern_ / cur_k\n",
    "                sparse_patterns.append(sparse_pattern_)\n",
    "            sparse_patterns = torch.stack(sparse_patterns, dim=-1) # [N,D,M]\n",
    "            # Reg Type1:  encourage the kernel to lean to map with more matching neighbors\n",
    "            temp_ = 0.2\n",
    "            eps = 1.e-3\n",
    "            sparse_patterns = F.softmax((F.softmax((sparse_patterns+eps)/temp_, dim=1)+eps)/temp_, dim=-1)  # [N. vec-dim. M] \n",
    "            self.register_buffer(\"sparse_patterns\",sparse_patterns)\n",
    "\n",
    "            return choice*self.sparse_patterns\n",
    "        else:\n",
    "            # formula 2: MultiScale Estimation of how sparse a point is \n",
    "            # apply softmax in the normalized N points dimension\n",
    "            # calc the relative sparsity distance to many centers as regs\n",
    "            raise NotImplementedError\n",
    "\n",
    "    def schedule_update(self, iter_=None):\n",
    "        '''\n",
    "        some schedulable params\n",
    "        '''\n",
    "        # ======= the temp annealing for choice =============\n",
    "        \n",
    "        # exponential temp annealing, best results in v3, now used in v5\n",
    "        self.temp = (self.temp_)**(1-iter_) # start from the temp, end with 0.00000....\n",
    "        \n",
    "        # linear temp annealing - had worse results than upper annealing - this was v4 test\n",
    "        # as in: https://www.researchgate.net/publication/337856246_Dynamic_Convolution_Attention_over_Convolution_Kernels\n",
    "        #if (self.temp_ - 3 * iter_ >= 0):\n",
    "        #    self.temp = self.temp_ - 3 * iter_ # start from the temp, end with 1\n",
    "        #else:\n",
    "        #    self.temp = 1\n",
    "\n",
    "        if self.skip_choice == True and iter_> 0.1:\n",
    "            self.skip_choice = False\n",
    "            print('SkipChoice Warmup Done, Start training choice qk')\n",
    "\n",
    "        if self.skip_choice == False and not hasattr(self, \"q\"):\n",
    "            self.q = nn.Sequential(\n",
    "                ME.MinkowskiConvolution(self.planes, self.M, kernel_size=3,dimension=3),\n",
    "                ME.MinkowskiBatchNorm(self.M),\n",
    "                    )\n",
    "            self.q.to(self.codebook[0][0].kernel.device)\n",
    "\n",
    "        pass\n",
    "\n",
    "        # ========= Temperature Annealing ==============\n",
    "        #if not hasattr(self, 'temp0'):\n",
    "        #    self.temp0 = self.temp\n",
    "\n",
    "        #self.temp = self.temp0*(0.01)**(iter_)\n",
    "\n",
    "    def forward(self, x, iter_=None, aux=None):\n",
    "        '''\n",
    "        For each dilation(D=3), different d have different kernel shape and different Ks, e.g., cube-shape kernel has k=27, cross-shaped has k=7\n",
    "        1st do qk projection: [N, dim, K]  ()\n",
    "                - conv: directly use conv neighbor aggregation(extra params), output: [N, H]\n",
    "                - pairwise: use linear mapping, then gather neighbor & dotproduct. output: [N, H, K] -> [N, H]\n",
    "        2nd: q_ dot product with Codebook(M set of conv weights): [N, H, M] -> [N, dim, M], the apply softmax to get choice of [D, M]\n",
    "        3rd: use choice: [D, M] to aggregate M codebook elements(channel-wise convs) for each point, then apply the coedbook(through channel-wise conv on value)\n",
    "        '''\n",
    "        self.register_buffer('coord_map', x.C)\n",
    "        self.schedule_update(iter_)\n",
    "\n",
    "        # align the channel for the decoder that concat the input\n",
    "        if self.planes != self.inplanes:\n",
    "            res = self.downsample(x)\n",
    "            x = self.linear_top(x)\n",
    "        else:\n",
    "            res = x\n",
    "\n",
    "        # generate the value\n",
    "        v_ = self.v(x)\n",
    "\n",
    "        # generate the qk\n",
    "        if self.skip_choice:\n",
    "            pass\n",
    "        else:  # no skip choice\n",
    "            if self.qk_type == 'conv':\n",
    "                if not self.CUSTOM_KERNEL:\n",
    "                    q_ = self.q(x)\n",
    "                    q_f = self.expand_dim(q_.F)\n",
    "                    q_= ME.SparseTensor(features=q_f, coordinate_map_key=q_.coordinate_map_key, coordinate_manager=q_.coordinate_manager) # [N, dim]\n",
    "                    N, dim = q_.F.shape\n",
    "                    qs = [q_]*self.D\n",
    "                else:\n",
    "                    qs = []\n",
    "                    for _ in range(self.D):\n",
    "                        q_ = self.q[_](x)\n",
    "                        q_f =self.expand_dim(q_.F)\n",
    "                        qs.append(\n",
    "                            ME.SparseTensor(features=q_f, coordinate_map_key=q_.coordinate_map_key, coordinate_manager=q_.coordinate_manager) # [N, dim]\n",
    "                                )\n",
    "                        N, dim = q_f.shape\n",
    "\n",
    "                # get dot-product of codebook-weight & q_\n",
    "                choice = []\n",
    "                out = []\n",
    "                for _ in range(self.D):\n",
    "                    self.codebook[_][0].kernel.requires_grad = False   # detach the grad from choice to codebook elements\n",
    "                    choice_ = self.codebook[_](qs[_])\n",
    "                    choice.append(choice_.F.reshape(\n",
    "                        [choice_.shape[0], self.M, self.planes*self.expansion // self.M]\n",
    "                            ).sum(-1)\n",
    "                        )\n",
    "                choice = torch.stack(choice, dim=-1)\n",
    "                eps = 1.e-3\n",
    "\n",
    "                if self.D > 1: # if M==1, skip softmax since there is only 1 value\n",
    "                    choice = F.softmax((choice)/self.temp, dim=-1) # [N, vec_dim, M] \n",
    "                else:\n",
    "                    pass\n",
    "\n",
    "                # attn_map = torch.stack([self.codebook[_][0].kernel for _ in range(self.D) ], dim=0) # [M. K], in some case(CUSTOM_KERNEL)\n",
    "                attn_map = torch.cat([self.codebook[_][0].kernel for _ in range(self.D)],dim=0) # [M. K]\n",
    "                self.register_buffer('attn_map', attn_map)\n",
    "                self.register_buffer('choice_map', choice)\n",
    "\n",
    "            elif self.qk_type == 'pairwise':\n",
    "\n",
    "                q_ = self.q(x)\n",
    "                q_f = q_.F\n",
    "                N, _ = q_.F.shape\n",
    "\n",
    "                choices = []\n",
    "                for i_m, kg in enumerate(self.kgargs):  # iter over M\n",
    "                    if 'dimension' in kg.keys():\n",
    "                        del kg['dimension']\n",
    "                    neis_d = q_.coordinate_manager.get_kernel_map(q_.coordinate_map_key,\n",
    "                                                                    q_.coordinate_map_key,\n",
    "                                                                    **kg\n",
    "                                                                        )\n",
    "                    choice = []\n",
    "                    for k_ in range(len(neis_d.keys())):\n",
    "                        if not k_ in neis_d.keys():\n",
    "                            continue\n",
    "                        neis_ = torch.gather(q_.F, dim=0, index=neis_d[k_][0].reshape(-1,1).expand(-1,self.M).long())\n",
    "                        neis = torch.zeros(N,self.M, device=q_.F.device)  # DEBUG: not sure if needs decalre every time\n",
    "                        neis.scatter_(dim=0, index=neis_d[k_][1].reshape(-1,1).expand(-1,self.M).long(), src=neis_)\n",
    "\n",
    "                        sparse_mask_cur_k = (neis.abs().sum(-1) > 0).float()\n",
    "                        neis = neis*(q_.F*sparse_mask_cur_k.unsqueeze(-1).expand(-1, self.M))\n",
    "                        neis = neis*sparse_mask_cur_k.unsqueeze(-1).expand(-1, self.M)\n",
    "\n",
    "                        out_cur_k = self.expand_dim(neis)*self.codebook[i_m][0].kernel[k_].unsqueeze(0)\n",
    "                        out_cur_k = out_cur_k.sum(1)  # [N]\n",
    "                        choice.append(out_cur_k)\n",
    "\n",
    "                    choice = torch.stack(choice, dim=-1)  # [N,K]\n",
    "                    choice = F.softmax(choice/self.temp, dim=-1).sum(-1)\n",
    "                    choices.append(choice) # [N]\n",
    "\n",
    "                choices = torch.stack(choices, dim=-1)\n",
    "                choices = F.softmax(choices/self.temp, dim=-1)   # [N,M]\n",
    "                choice = choices.unsqueeze(1).expand(-1, self.M, -1) # [N, dim, M]\n",
    "                self.register_buffer('choice_map', choices)\n",
    "\n",
    "            if self.sparse_pattern_reg:\n",
    "                choice = self.get_sparse_pattern(x, choice)\n",
    "\n",
    "        if self.skip_choice:\n",
    "            N, dim = v_.shape\n",
    "            out = []\n",
    "            for _ in range(self.D):\n",
    "                self.codebook[_][0].kernel.requires_grad = True\n",
    "                out_ = self.codebook[_](v_)\n",
    "                out.append(out_.F)\n",
    "            out = torch.stack(out, dim=-1)\n",
    "            out = out.sum(-1)\n",
    "\n",
    "        elif self.top_k_choice:\n",
    "            assert self.M == 1 # same point use the same choice \n",
    "            out = torch.zeros([N,dim,self.top_k_choice], device=x.device)\n",
    "            choice_topk = torch.topk(choice, self.top_k_choice, dim=-1)[0] # shape [N,dim]\n",
    "            choice_topk_idx = torch.topk(choice, self.top_k_choice, dim=-1)[1][:,0,:]  # shape [N]\n",
    "            for _ in range(self.D):\n",
    "                self.codebook[_][0].kernel.requires_grad = True\n",
    "                # DEV: split points for different choice\n",
    "                # however, if choice has the channle freedom\n",
    "                # could not handle\n",
    "                cur_out_ = self.codebook[_](v_) # the conv\n",
    "                for top_ in range(self.top_k_choice):\n",
    "                    choice_idx = torch.where(choice_topk_idx[:,top_] == _)[0]\n",
    "                    # cur_v_ = v_.features_at_coordinates(v_.C[choice_idx,:].float())\n",
    "                    if len(choice_idx) > 1:\n",
    "                        # cur_v_ = ME.SparseTensor(\n",
    "                                # features=v_.F[choice_idx,:],\n",
    "                                # coordinates=v_.C[choice_idx,:],\n",
    "                                # coordinate_map_key=v_.coordinate_map_key,\n",
    "                                # coordinate_manager=v_.coordinate_manager\n",
    "                                # )\n",
    "                        out[:,:,top_].scatter_(\n",
    "                                src=cur_out_.F[choice_idx,:]*choice_topk[choice_idx,:,top_],\n",
    "                                index=choice_idx.unsqueeze(-1).repeat(1,dim),\n",
    "                                dim=0)\n",
    "                    else:\n",
    "                        pass\n",
    "            out = out.sum(-1)\n",
    "        else:\n",
    "            # normal-case: apply the attn_weight aggregation with the channelwiseConvolution\n",
    "            out = torch.zeros([N, self.planes*self.expansion], device=v_.device)\n",
    "            for _ in range(self.D):\n",
    "                self.codebook[_][0].kernel.requires_grad = True\n",
    "                out_ = self.codebook[_](v_)\n",
    "                out += out_.F*self.expand_dim(choice[:,:,_])\n",
    "            out = out.reshape([N, self.planes*self.expansion])\n",
    "\n",
    "        out = ME.SparseTensor(features=out, coordinate_map_key=x.coordinate_map_key, coordinate_manager=x.coordinate_manager)\n",
    "        out = self.out_bn_relu(out)\n",
    "        out = out + res\n",
    "\n",
    "        return out\n",
    "\n",
    "        # argumentation\n",
    "        # 1. rechenaufwand -> wir haben encoder/decoder -> wichtigste attn blöcke sind in der mitte weil features most condensed\n",
    "        # warum im decoder zuerst streichen: einzelpunktinformationen im encoder noch unverzerrt(er) \n",
    "    \n",
    "        # 2. warum zwischen jedem block -> attention is all you need paper - wir haben conv blöcke, machine translation hat feed forward blöcke \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d340482d",
   "metadata": {},
   "source": [
    "## AttentiveCylinder3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ef6bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################################################################################\n",
    "# Net Architecture - Second Part                                                         #\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",
    "        self.head_size = 8\n",
    "\n",
    "        sparse_shape = np.array(output_shape)\n",
    "        # sparse_shape[0] = 11\n",
    "        print(\"sparse shape:\" + str(sparse_shape))\n",
    "        # cylindrical partition splits these point clouds into 3D representation\n",
    "        # with the size = 480 × 360× 32, where three dimensions indicate the radius,\n",
    "        # angle and height, respectively.\n",
    "        self.sparse_shape = sparse_shape\n",
    "\n",
    "        self.downCntx = ResContextBlock(num_input_features, init_size, indice_key=\"pre\")\n",
    "        self.downAttn2Block = CodedVTRBlock(init_size, init_size)\n",
    "        \n",
    "        # DOWN BLOCKS\n",
    "        self.resBlock2 = ResBlock(init_size, 2 * init_size, 0.2, height_pooling=True, indice_key=\"down2\")\n",
    "        self.downAttn3Block = CodedVTRBlock(2 * init_size, 2 * init_size)\n",
    "        self.resBlock3 = ResBlock(2 * init_size, 4 * init_size, 0.2, height_pooling=True, indice_key=\"down3\")\n",
    "        self.downAttn4Block = CodedVTRBlock(4 * init_size, 4 * init_size)\n",
    "        self.resBlock4 = ResBlock(4 * init_size, 8 * init_size, 0.2, pooling=True, height_pooling=False, indice_key=\"down4\")\n",
    "        self.downAttn5Block = CodedVTRBlock(8 * init_size, 8 * init_size)\n",
    "        self.resBlock5 = ResBlock(8 * init_size, 16 * init_size, 0.2, pooling=True, height_pooling=False, indice_key=\"down5\")\n",
    "        \n",
    "        # UP BLOCKS with CodedVTR Attention Block    \n",
    "        self.upAttn0Block = CodedVTRBlock(16 * init_size, 16 * init_size)\n",
    "        self.upBlock0 = UpBlock(16 * init_size, 16 * init_size, indice_key=\"up0\", up_key=\"down5\")    \n",
    "        self.upAttn1Block = CodedVTRBlock(16 * init_size, 16 * init_size)\n",
    "        self.upBlock1 = UpBlock(16 * init_size, 8 * init_size, indice_key=\"up1\", up_key=\"down4\")  \n",
    "        self.upAttn2Block = CodedVTRBlock(8 * init_size, 8 * init_size)\n",
    "        self.upBlock2 = UpBlock(8 * init_size, 4 * init_size, indice_key=\"up2\", up_key=\"down3\") \n",
    "        self.upAttn3Block = CodedVTRBlock(4 * init_size, 4 * init_size)\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, iter):\n",
    "        # x = x.contiguous()\n",
    "        coors = coors.int()\n",
    "        # import pdb\n",
    "        # pdb.set_trace()\n",
    "        sparseconv_features = spconv.SparseConvTensor(voxel_features, coors, self.sparse_shape,\n",
    "                                      batch_size)\n",
    "        \n",
    "        #################################\n",
    "        # Down Context Block            #\n",
    "        #################################\n",
    "        downcntx = self.downCntx(sparseconv_features)\n",
    "        \n",
    "        #################################\n",
    "        # ENCODER BLOCKS                #\n",
    "        #################################\n",
    "        downAttn1 = self.downAttn2Block(ME.SparseTensor(features=downcntx.features, coordinates=downcntx.indices), iter_=iter)\n",
    "        downcntx.replace_feature(downAttn1.F)\n",
    "        \n",
    "        down1c, down1b = self.resBlock2(downcntx)\n",
    "        \n",
    "        downAttn2 = self.downAttn3Block(ME.SparseTensor(features=down1c.features, coordinates=down1c.indices), iter_=iter)\n",
    "        down1c.replace_feature(downAttn2.F)\n",
    "        \n",
    "        down2c, down2b = self.resBlock3(down1c)\n",
    "        \n",
    "        downAttn3 = self.downAttn4Block(ME.SparseTensor(features=down2c.features, coordinates=down2c.indices), iter_=iter)\n",
    "        down2c.replace_feature(downAttn3.F)\n",
    "        \n",
    "        down3c, down3b = self.resBlock4(down2c)\n",
    "        \n",
    "        downAttn4 = self.downAttn5Block(ME.SparseTensor(features=down3c.features, coordinates=down3c.indices), iter_=iter)\n",
    "        down3c.replace_feature(downAttn4.F)\n",
    "        \n",
    "        down4c, down4b = self.resBlock5(down3c)     \n",
    "        \n",
    "        #################################\n",
    "        # DECODER BLOCKS                #\n",
    "        #################################\n",
    "        upAttn4 = self.upAttn0Block(ME.SparseTensor(features=down4c.features, coordinates=down4c.indices), iter_=iter)\n",
    "        down4c.replace_feature(upAttn4.F)        \n",
    "\n",
    "        up4e = self.upBlock0(down4c, down4b)\n",
    "                \n",
    "        upAttn3 = self.upAttn1Block(ME.SparseTensor(features=up4e.features, coordinates=up4e.indices), iter_=iter)\n",
    "        up4e.replace_feature(upAttn3.F)\n",
    "        \n",
    "        up3e = self.upBlock1(up4e, down3b)\n",
    "               \n",
    "        upAttn2 = self.upAttn2Block(ME.SparseTensor(features=up3e.features, coordinates=up3e.indices), iter_=iter)\n",
    "        up3e.replace_feature(upAttn2.F)\n",
    "\n",
    "        up2e = self.upBlock2(up3e, down2b) \n",
    "        \n",
    "        upAttn1 = self.upAttn3Block(ME.SparseTensor(features=up2e.features, coordinates=up2e.indices), iter_=iter)\n",
    "        up2e.replace_feature(upAttn1.F)\n",
    "        \n",
    "        up1e = self.upBlock3(up2e, down1b)\n",
    "        \n",
    "        #################################\n",
    "        # DDCM BLOCK                    #\n",
    "        #################################\n",
    "        up0e = self.ReconNet(up1e)\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": "a33bccb2-456a-42e8-9542-ba2e143987db",
   "metadata": {},
   "source": [
    "## Cylinder Feature Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "da70259b-b4c1-45ea-ab49-bafa8dd531c0",
   "metadata": {},
   "source": [
    "## Loss Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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",
    "        train_data_path = train_dataloader_config[\"data_path\"]\n",
    "        val_data_path = val_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(train_data_path, imageset=train_imageset,\n",
    "                                    return_ref=train_ref, label_mapping=label_mapping, nusc=nusc)\n",
    "        val_pt_dataset = SemKITTI(val_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": null,
   "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",
    "        # Net Architecture - First Part\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",
    "        # Net Architecture - Second Part\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",
    "\n",
    "        # Put everything together, call forward and return the model\n",
    "        model = get_model_class_c3d(model_config[\"model_architecture\"])(\n",
    "            cylin_model=cy_fea_net,                            # Net Architecture - First Part\n",
    "            segmentator_spconv=cylinder_3d_spconv_seg,         # Net Architecture - Second Part\n",
    "            sparse_shape=output_shape                          \n",
    "        )\n",
    "\n",
    "        return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3398392",
   "metadata": {},
   "source": [
    "## Build Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f6744ec",
   "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(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": "a70e592b-11a6-491e-930f-c86c64520e8c",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d086ba0b-972c-44af-9c67-ce0b89504c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluation(config_path, demo_folder, demo_label_folder, save_folder):\n",
    "    pytorch_device = torch.device('cuda:0')\n",
    "    configs = load_config_data(config_path)\n",
    "\n",
    "    dataset_config = configs['dataset_params']\n",
    "    train_dataloader_config = configs['train_data_loader']\n",
    "    val_dataloader_config = configs['val_data_loader']\n",
    "\n",
    "    val_batch_size = val_dataloader_config['batch_size']\n",
    "    train_batch_size = train_dataloader_config['batch_size']\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",
    "\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",
    "\n",
    "    model_load_path = train_hypers['model_load_path']\n",
    "    model_save_path = train_hypers['model_save_path']\n",
    "    model_save_path_early_stop = train_hypers['model_save_path_early_stop']\n",
    "    model_save_path_best_val_miou = train_hypers['model_save_path_best_val_miou']\n",
    "\n",
    "    SemKITTI_label_name = get_SemKITTI_label_name(dataset_config[\"label_mapping\"])\n",
    "    \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",
    "    \n",
    "    my_model = model_builder.build(model_config)\n",
    "    \n",
    "    if os.path.exists(model_load_path):\n",
    "        my_model = load_checkpoint(model_load_path, my_model)\n",
    "        print(\"pretrained_checkpoint loaded: \" + model_load_path)\n",
    "    else:\n",
    "        print(\"no pretrained_checkpoint loaded.\")\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, iter=0)\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",
    "                # watch out, sbld starts with \n",
    "                outputPath = save_dir + str(i_iter_demo+1).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": [
    "# Evaluation / Label Prediction Settings # 09 14 25\n",
    "#config_path = \"config/semantickitti.yaml\"\n",
    "config_path = \"config/sbld.yaml\"\n",
    "demo_folder = \"../eval_set/28/velodyne/\"\n",
    "save_folder = \"../eval_set/28/predictions_attentivecylinder3d\"\n",
    "demo_label_folder = \"../eval_set/28/labels/\"\n",
    "\n",
    "# Evaluation\n",
    "evaluation(config_path, demo_folder, demo_label_folder, save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aef4803",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}