0311_predict_tlcnetU_process_wholeimg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TLCUNET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## predict whole image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "import os\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import tifffile as tif\n",
    "\n",
    "from ptsemseg.metrics import heightacc\n",
    "from ptsemseg.models import TLCNetU\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import tifffile as tif\n",
    "from os.path import join\n",
    "import math\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: spectral in d:\\code\\msws\\venv\\lib\\site-packages (0.23.1)\n",
      "Requirement already satisfied: numpy in d:\\code\\msws\\venv\\lib\\site-packages (from spectral) (1.25.2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip available: 22.3.1 -> 23.3.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "!pip install spectral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup device\n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rasterio as rio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from osgeo import gdal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup Model\n",
    "model = TLCNetU(n_classes=1).to(device)\n",
    "if torch.cuda.device_count() > 1:\n",
    "    model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=> loading checkpoint 'runs\\tlcnetu_zy3bh\\finetune.tar'\n",
      "=> loaded checkpoint 'runs\\tlcnetu_zy3bh\\finetune.tar' (epoch 298)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TLCNetU(\n",
       "  (uencoder1): Uencoder(\n",
       "    (conv1): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv2): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv3): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv4): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (center): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (uencoder2): Uencoder(\n",
       "    (conv1): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(4, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv2): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv3): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (conv4): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "    (maxpool4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (center): unetConv2(\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (2): ReLU()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (udecoder1): UdecoderC(\n",
       "    (up_concat4): unetUpC(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(512, 128, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat3): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat2): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat1): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(32, 16, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (final): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (udecoder2): Udecoder(\n",
       "    (up_concat4): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat3): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat2): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat1): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(32, 16, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (final): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (udecoder3): Udecoder(\n",
       "    (up_concat4): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat3): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat2): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (up_concat1): unetUp(\n",
       "      (conv): unetConv2(\n",
       "        (conv1): Sequential(\n",
       "          (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "        (conv2): Sequential(\n",
       "          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU()\n",
       "        )\n",
       "      )\n",
       "      (up): ConvTranspose2d(32, 16, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (final): Conv2d(16, 2, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (final): Conv2d(4, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#resume = cfg[\"training\"][\"resume\"]\n",
    "resume = r'runs\\tlcnetu_zy3bh\\finetune.tar'\n",
    "if os.path.isfile(resume):\n",
    "    print(\"=> loading checkpoint '{}'\".format(resume))\n",
    "    checkpoint = torch.load(resume)\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    # optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "    print(\"=> loaded checkpoint '{}' (epoch {})\"\n",
    "             .format(resume, checkpoint['epoch']))\n",
    "else:\n",
    "    print(\"=> no checkpoint found at resume\")\n",
    "    print(\"=> Will start from scratch.\")\n",
    "\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_whole_image(model, image, r, c, grid=400):\n",
    "    '''\n",
    "    image: n,r,c,b  where n = 1\n",
    "    model: FCN\n",
    "    non-overlap\n",
    "    '''\n",
    "    #grid=400\n",
    "    n, b, rows, cols = image.shape\n",
    "    # rows=math.ceil(r/grid)*grid\n",
    "    # cols=math.ceil(c/grid)*grid\n",
    "    # image_= np.pad(image,((0,0), (0,0), (0,rows-r), (0,cols-c)),'symmetric')\n",
    "    # weight = np.ones((rows, cols))\n",
    "    res = np.zeros((rows, cols),dtype=np.float32)\n",
    "    seg = np.zeros((rows, cols), dtype=np.uint8)\n",
    "    num_patch= len(range(0,rows,grid))*len(range(0,cols,grid))\n",
    "    print('num of patch is',num_patch)\n",
    "    k=0\n",
    "    for i in range(0,rows, grid):\n",
    "        for j in range(0, cols, grid):          \n",
    "            patch = image[0:,0:,i:i+grid,j:j+grid].astype('float32')\n",
    "            if np.max(patch.flatten())<=10e-8:\n",
    "                continue\n",
    "            start=time.time()\n",
    "            patch = torch.from_numpy(patch).float()\n",
    "            \n",
    "            # normalization \n",
    "            patch = patch/1000.0\n",
    "\n",
    "            pred = model(patch.to(device))\n",
    "            pred0 = pred[0].cpu().detach().numpy() # height\n",
    "            pred1 = pred[1].cpu().detach().numpy() # seg\n",
    "            res[i:i+grid,j:j+grid] = np.squeeze(pred0)\n",
    "            seg[i:i+grid,j:j+grid] = np.argmax(np.squeeze(pred1), axis=0) # argmax\n",
    "            end=time.time()\n",
    "            k=k+1\n",
    "            #print('patch [%d/%d] time elapse:%.3f'%(k,num_patch,(end-start)))\n",
    "    res = res[0:r,0:c].astype(np.float32)\n",
    "    seg = seg[0:r,0:c].astype(np.uint8)\n",
    "    return res, seg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## batch processing "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read all city names\n",
    "def read_filepath(filepath, filename):\n",
    "    filelist=list()\n",
    "    for root, dirs, files in os.walk(filepath):\n",
    "        for name in files:\n",
    "            if name.endswith(filename):\n",
    "                filelist.append(join(root, name))\n",
    "    filelist.sort()\n",
    "    return filelist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "filelistnew = [r'Z:\\yinxcao\\hongkong\\mux_quacss1.tif']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid =400"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loadenvi(file, band=1):\n",
    "    dataset = gdal.Open(file)\n",
    "    band = dataset.GetRasterBand(band)\n",
    "    data = band.ReadAsArray()\n",
    "    dataset = None\n",
    "    band = None\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(16800, 20800)\n",
      "uint16\n"
     ]
    }
   ],
   "source": [
    "data = loadenvi(filelistnew[0], band=1)\n",
    "print(data.shape)\n",
    "print(data.dtype)\n",
    "data = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "process: %s D:\\data\\zy3height\\ningbo\\mux_quacs\n",
      "num of patch is 2184\n",
      "success: D:\\data\\zy3height\\ningbo\\mux_quacs, time is 40.753537\n"
     ]
    }
   ],
   "source": [
    "for file in filelistnew[:1]:\n",
    "    # filepath\n",
    "    idirname=os.path.dirname(file)\n",
    "    predpath = join(idirname, 'pred')\n",
    "    respath=join(predpath, 'predtlcnetu_200.tif') # model tlcnetu, epoch 200\n",
    "    segpath=join(predpath, 'predtlcnetu_200_seg.tif')\n",
    "    if os.path.exists(respath):\n",
    "        print('file: %s already exist, then skip'%respath)\n",
    "        continue\n",
    "    if not os.path.exists(predpath):\n",
    "        print('mkdir %s'%predpath)\n",
    "        os.mkdir(predpath)\n",
    "        \n",
    "    print('process: %s', file)\n",
    "    nadpath=join(idirname,'nads')\n",
    "    fwdpath=join(idirname,'fwds')\n",
    "    bwdpath=join(idirname,'bwds')\n",
    "    # 1.read image\n",
    "    band1 = loadenvi(file, band=1)\n",
    "    r,c =band1.shape[:2]\n",
    "    img=np.zeros((r,c, 7), dtype='uint16')\n",
    "    for i in range(4):\n",
    "        img[:, :, i] = loadenvi(file, i+1)\n",
    "    img[:, :, 4]=loadenvi(nadpath)\n",
    "    img[:, :, 5]=loadenvi(fwdpath)\n",
    "    img[:, :, 6]=loadenvi(bwdpath)\n",
    "    img = img.transpose(2, 0, 1) # C H W\n",
    "    \n",
    "    # img=img/10000\n",
    "    img = np.expand_dims(img, axis=0) # 1 C H W\n",
    "    \n",
    "    rows=math.ceil(r/grid)*grid\n",
    "    cols=math.ceil(c/grid)*grid\n",
    "    img = np.pad(img,((0,0), (0,0), (0,rows-r), (0,cols-c)),'symmetric')\n",
    "\n",
    "    # 2. predict\n",
    "    starttime=time.time()\n",
    "    res=predict_whole_image(model, img, r, c,  grid=400)\n",
    "    endtime=time.time()\n",
    "    print('success: %s, time is %f'%(file, endtime-starttime))\n",
    "\n",
    "    img = 0 # release \n",
    "    \n",
    "    # 3. export\n",
    "    reffile = file+'8.tif'\n",
    "    rastermeta = rio.open(reffile).profile\n",
    "    rastermeta.update(dtype=res[0].dtype, count=1, compress='lzw')\n",
    "    with rio.open(respath, mode=\"w\", **rastermeta) as dst:\n",
    "        dst.write(res[0],1)\n",
    "    rastermeta.update(dtype=res[1].dtype, count=1, compress='lzw')\n",
    "    with rio.open(segpath, mode=\"w\", **rastermeta) as dst:\n",
    "        dst.write(res[1],1)\n",
    "#     tif.imsave(respath, res[0]) # building height\n",
    "#     tif.imsave(segpath, res[1]) # building segmentation"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}