ISPRS_4_layer_code.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r5ZXY2XcC3Pm"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict\n",
        "from torch import Tensor\n",
        "import re\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.checkpoint as cp\n",
        "from collections import OrderedDict\n",
        "from torch import Tensor\n",
        "from typing import Type, Any, Callable, Union, List, Optional\n",
        "import glob\n",
        "import os\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from torch import optim\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
        "from torch.utils.data import Subset,DataLoader\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision\n",
        "import random\n",
        "from google.colab import files\n",
        "from sklearn.metrics import  confusion_matrix\n",
        "from sklearn.model_selection import ShuffleSplit\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "from scipy.ndimage import distance_transform_edt\n",
        "from torch.autograd import Variable\n",
        "import skimage.segmentation\n",
        "import skimage.io\n",
        "import skimage\n",
        "from scipy.optimize import linear_sum_assignment\n",
        "import skimage.segmentation\n",
        "import matplotlib.pyplot as plt\n",
        "import skimage.io\n",
        "import skimage.segmentation\n",
        "from skimage import feature\n",
        "from skimage import filters\n",
        "import copy\n",
        "import torchvision\n",
        "from collections import OrderedDict\n",
        "import math\n",
        "import imageio\n",
        "from skimage.color import rgb2gray\n",
        "\n",
        "from torchsummary import summary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gf1GtdJD9e_1"
      },
      "outputs": [],
      "source": [
        "def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:\n",
        "    \"\"\"3x3 convolution with padding\"\"\"\n",
        "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
        "                     padding=dilation, groups=groups, bias=False, dilation=dilation)\n",
        "\n",
        "\n",
        "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:\n",
        "    \"\"\"1x1 convolution\"\"\"\n",
        "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class Upsample(nn.Module):\n",
        "    \"\"\" nn.Upsample is deprecated \"\"\"\n",
        "\n",
        "    def __init__(self, scale_factor, mode=\"bilinear\"):\n",
        "        super(Upsample, self).__init__()\n",
        "        self.scale_factor = scale_factor\n",
        "        self.mode = mode\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=True, recompute_scale_factor=True)\n",
        "        return x"
      ],
      "metadata": {
        "id": "n_evk9rlrTb-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Index Pooling\n",
        "class pool(nn.Module):\n",
        "  def __init__(self,channels):\n",
        "    super(pool, self).__init__()\n",
        "    self.channels=channels\n",
        "    self.weight1 = torch.zeros((channels,1,2,2)).cuda()\n",
        "    self.weight2 = torch.zeros((channels,1,2,2)).cuda()\n",
        "    self.weight3 = torch.zeros((channels,1,2,2)).cuda()\n",
        "    self.weight4 = torch.zeros((channels,1,2,2)).cuda()\n",
        "    self.weight1[:,:,0,0]=1\n",
        "    self.weight2[:,:,0,1]=1\n",
        "    self.weight3[:,:,1,0]=1\n",
        "    self.weight4[:,:,1,1]=1\n",
        "  def forward(self, x):\n",
        "    with torch.no_grad():\n",
        "         x1=F.conv2d(x, self.weight1,stride=2,groups=self.channels, bias=None)\n",
        "         x2=F.conv2d(x, self.weight2,stride=2,groups=self.channels, bias=None)\n",
        "         x3=F.conv2d(x, self.weight3,stride=2,groups=self.channels, bias=None)\n",
        "         x4=F.conv2d(x, self.weight4,stride=2,groups=self.channels, bias=None)\n",
        "    return x1,x2,x3,x4"
      ],
      "metadata": {
        "id": "8Rd0D6KdrVV2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#DAMIP Module\n",
        "class attn_pool(nn.Module):\n",
        "    def __init__(self,feature_channels):\n",
        "       super(attn_pool, self).__init__()\n",
        "       self.pool1=pool(feature_channels)\n",
        "       self.pool2=pool(1)\n",
        "       self.conv1 = nn.Conv2d(feature_channels*4, 2*feature_channels, kernel_size=1, stride=1, bias=False)\n",
        "       self.conv2 = nn.Conv2d(2*feature_channels,2*feature_channels,kernel_size=7,stride=1,padding=3)\n",
        "       self.bn2 = nn.BatchNorm2d(2*feature_channels)\n",
        "       self.a1=nn.Parameter(torch.Tensor(1))\n",
        "       self.a2=nn.Parameter(torch.Tensor(1))\n",
        "       self.a3=nn.Parameter(torch.Tensor(1))\n",
        "       self.a4=nn.Parameter(torch.Tensor(1))\n",
        "    def forward(self,map,feature):\n",
        "        feature1,feature2,feature3,feature4=self.pool1(feature)\n",
        "        map1,map2,map3,map4=self.pool2(map)\n",
        "\n",
        "        fm1 = self.a1*feature1 + feature1*map1\n",
        "        fm2 = self.a2*feature2 + feature2*map2\n",
        "        fm3 = self.a3*feature3 + feature3*map3\n",
        "        fm4 = self.a4*feature4 + feature4*map4\n",
        "\n",
        "        mat=torch.cat((fm1,fm2,fm3,fm4),1)\n",
        "        mat=self.conv1(mat)\n",
        "        mat=F.relu(self.bn2(self.conv2(mat)))\n",
        "        return mat"
      ],
      "metadata": {
        "id": "EFfvKxnLrXT7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# DPMG Module\n",
        "class dsup(nn.Module):\n",
        "  def __init__(self,input_channels):\n",
        "    super(dsup,self).__init__()\n",
        "    self.conv1=conv3x3(input_channels,input_channels//2)\n",
        "    self.bn1=nn.BatchNorm2d(input_channels//2)\n",
        "    self.conv2=conv3x3(input_channels//2,32)\n",
        "    self.conv3=nn.Conv2d(32,1,kernel_size=3,stride=1,padding=1)\n",
        "  def forward(self,x):\n",
        "    x=F.relu(self.bn1(self.conv1(x)))\n",
        "    x=self.conv2(x)\n",
        "    x=self.conv3(x)\n",
        "    return torch.sigmoid(x)"
      ],
      "metadata": {
        "id": "CYY4YvlgrZHA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Dilated Convolutional Block\n",
        "class conv_enc(nn.Module):\n",
        "  def __init__(self,in_channels,out_channels,dil):\n",
        "    super(conv_enc,self).__init__()\n",
        "    self.conv1 = conv1x1(in_channels,in_channels)\n",
        "    self.bn1 = nn.BatchNorm2d(in_channels)\n",
        "    self.conv2 = conv3x3(in_channels,in_channels,dilation=dil)\n",
        "    self.bn2 = nn.BatchNorm2d(in_channels)\n",
        "    self.conv3 = conv1x1(in_channels,out_channels)\n",
        "    self.bn3 = nn.BatchNorm2d(out_channels)\n",
        "    self.conv4 = conv3x3(out_channels,out_channels,dilation=dil)\n",
        "    self.bn4 = nn.BatchNorm2d(out_channels)\n",
        "    # input dimension matching\n",
        "    self.conv0 = conv1x1(in_channels,out_channels)\n",
        "  def forward(self,x):\n",
        "    identity = self.conv0(x)\n",
        "\n",
        "    x=F.relu(self.bn1(self.conv1(x)))\n",
        "    x=F.relu(self.bn2(self.conv2(x)))\n",
        "    x=F.relu(self.bn3(self.conv3(x)))\n",
        "    x=self.bn4(self.conv4(x))\n",
        "\n",
        "    return F.relu(x+identity)"
      ],
      "metadata": {
        "id": "q4NFYAvGra-Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Feature Extractor\n",
        "class enc(nn.Module):\n",
        "  def __init__(self,input_channels,output_channels,dil):\n",
        "    super(enc,self).__init__()\n",
        "    self.conv1 = conv_enc(input_channels,output_channels,dil)\n",
        "    self.conv2 = conv_enc(output_channels,output_channels,2*dil)\n",
        "    self.dp_sup = dsup(output_channels)     # DPMG Module\n",
        "\n",
        "  def forward(self,x):\n",
        "    x1 = self.conv1(x)\n",
        "    x1 = self.conv2(x1)\n",
        "    x1_out = self.dp_sup(x1)\n",
        "    return x1, x1_out"
      ],
      "metadata": {
        "id": "PGsTFmTrrczn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# DAMSCA Module\n",
        "class kqcbam(nn.Module):\n",
        "  def __init__(self,input_channels,scale_factor=2):\n",
        "    super(kqcbam,self).__init__()\n",
        "    self.conv1=nn.Conv2d(1,input_channels,kernel_size=1)\n",
        "    self.gap=nn.AdaptiveAvgPool2d((1,1))\n",
        "    self.upsample=Upsample(scale_factor)\n",
        "  def forward(self,map,feature):\n",
        "    f1=map*feature\n",
        "    map2=self.conv1(map)\n",
        "    map2=self.gap(map2)\n",
        "    f2=torch.sigmoid(map2)*feature\n",
        "    out=F.relu(f1+f2)\n",
        "    return self.upsample(out)"
      ],
      "metadata": {
        "id": "4y-BJiKzrew7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Decoder Module\n",
        "class decoder(nn.Module):\n",
        "  def __init__(self,input_channels):\n",
        "    super(decoder,self).__init__()\n",
        "    self.conv1=nn.ConvTranspose2d(input_channels,128,kernel_size=4,stride=2,padding=1)\n",
        "    self.bn1=nn.BatchNorm2d(128)\n",
        "    self.conv2=nn.Conv2d(128,64,3,stride=1,padding=1)\n",
        "    self.bn2=nn.BatchNorm2d(64)\n",
        "    self.conv3=nn.Conv2d(64,32,3,stride=1,padding=1)\n",
        "    self.conv_out=nn.Conv2d(32,1,3,stride=1,padding=1)\n",
        "  def forward(self,x):\n",
        "    x=F.relu(self.bn1(self.conv1(x)))\n",
        "    x=F.relu(self.bn2(self.conv2(x)))\n",
        "    x=self.conv3(x)\n",
        "    x=self.conv_out(x)\n",
        "    return torch.sigmoid(x)"
      ],
      "metadata": {
        "id": "OG0Omr22rg09"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# MSSDMPA-Net\n",
        "class dsmpnet(nn.Module):\n",
        "  def __init__(self,input_channels):\n",
        "    super(dsmpnet,self).__init__()\n",
        "    self.conv1=nn.Conv2d(input_channels,64,kernel_size=7,stride=2,padding=3)\n",
        "    self.bn1=nn.BatchNorm2d(64)\n",
        "\n",
        "    self.pool1=attn_pool(64)\n",
        "    self.pool2=attn_pool(128)\n",
        "    self.pool3=attn_pool(256)\n",
        "\n",
        "    self.path1=enc(64,64,1)\n",
        "    self.path2=enc(128,128,2)\n",
        "    self.path3=enc(256,256,3)\n",
        "    self.path4=enc(512,512,4)\n",
        "\n",
        "    self.cbm1=kqcbam(64,1)\n",
        "    self.cbm2=kqcbam(128,2)\n",
        "    self.cbm3=kqcbam(256,4)\n",
        "    self.cbm4=kqcbam(512,8)\n",
        "\n",
        "    self.decoder=decoder(448)\n",
        "    self.decoder=decoder(960)\n",
        "\n",
        "  def forward(self,x):\n",
        "    x=F.relu(self.bn1(self.conv1(x)))\n",
        "    x1,x1_out=self.path1(x)\n",
        "    x=self.pool1(x1_out,x)\n",
        "    x2,x2_out=self.path2(x)\n",
        "    x=self.pool2(x2_out,x)\n",
        "    x3,x3_out=self.path3(x)\n",
        "    x=self.pool3(x3_out,x)\n",
        "    x4,x4_out=self.path4(x)\n",
        "\n",
        "    x1=self.cbm1(x1_out,x1)\n",
        "    x2=self.cbm2(x2_out,x2)\n",
        "    x3=self.cbm3(x3_out,x3)\n",
        "    x4=self.cbm4(x4_out,x4)\n",
        "    x_out=torch.cat((x1,x2,x3,x4),1)\n",
        "    x_out=self.decoder(x_out)\n",
        "    return x_out,x1_out,x2_out,x3_out,x4_out"
      ],
      "metadata": {
        "id": "fQU9GW0Frihl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PVvLrQfHb0pN"
      },
      "outputs": [],
      "source": [
        "model=dsmpnet(3).cuda()"
      ]
    }
  ]
}