LocUNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "anonymous-float",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, utils, datasets, models\n",
    "from scipy import ndimage\n",
    "import scipy.io\n",
    "# Ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "from lib import loader, modules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "joint-dominican",
   "metadata": {},
   "source": [
    "## Create results directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "encouraging-plenty",
   "metadata": {},
   "outputs": [],
   "source": [
    "try: \n",
    "    os.mkdir('results') \n",
    "except OSError as error: \n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conceptual-hello",
   "metadata": {},
   "source": [
    "## Simulation type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "focused-north",
   "metadata": {},
   "outputs": [],
   "source": [
    "simName = \"DPM\" # Options: DPM, ZSDPMtoIRT2, DPMtoIRT2, DPMcars, IRT2carsCDPM, IRT2carsCDPMtoIRT, \n",
    "                # from top to bottom in Table II of the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conditional-virtue",
   "metadata": {},
   "source": [
    "## Log file to save the best validation loss among the models of 50 training epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "requested-isaac",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results/' + simName + 'Log.txt', 'w') as f:\n",
    "    print('Training Accuracy', file=f)\n",
    "    print('-' * 20, file=f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "reserved-sleeping",
   "metadata": {},
   "source": [
    "## Loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compressed-genre",
   "metadata": {},
   "outputs": [],
   "source": [
    "Loc_train = loader.locDL(phase=\"train\",dir_dataset=\"dataset/\",cityMap=\"true\",carsMap=\"false\",simulation=simName,TxMaps=\"true\")\n",
    "Loc_val = loader.locDL(phase=\"val\",dir_dataset=\"dataset/\",cityMap=\"true\",carsMap=\"false\",simulation=simName,TxMaps=\"true\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acute-control",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloaders = {\n",
    "    'train': DataLoader(Loc_train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True),\n",
    "    'val': DataLoader(Loc_val, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "artistic-chamber",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "automatic-house",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_default_dtype(torch.float32)\n",
    "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "batch_size = 15\n",
    "\n",
    "inp = 16\n",
    "model = modules.LocUNet(inputs=inp)\n",
    "model.cuda()\n",
    "\n",
    "\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import time\n",
    "from datetime import datetime\n",
    "import copy\n",
    "from collections import defaultdict\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "\n",
    "def my_loss(output, target):\n",
    "    loss = torch.sum((output - target)**2,1)\n",
    "    loss = torch.sqrt(loss)\n",
    "    loss = torch.mean(loss)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def calc_loss_dense(pred, target, metrics):\n",
    "    loss = my_loss(pred, target)# *256*256\n",
    "    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)\n",
    "    return loss\n",
    "\n",
    "\n",
    "\n",
    "def print_metrics(metrics, epoch_samples, phase):\n",
    "    outputs1 = []\n",
    "    for k in metrics.keys():\n",
    "        outputs1.append(\"{}: {:4f}\".format(k, metrics[k] / epoch_samples))\n",
    "    with open('results/Log.txt', 'a') as f:\n",
    "        print(\"{}: {}\".format(phase, \", \".join(outputs1)), file=f)\n",
    "\n",
    "def train_model(model, optimizer, scheduler, num_epochs=50):\n",
    "    \n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_loss = 1e10\n",
    "    for epoch in range(num_epochs):\n",
    "        with open('results/Log.txt', 'a') as f:\n",
    "            print('Epoch {}/{}'.format(epoch, num_epochs - 1), file=f)\n",
    "            print('-' * 10, file=f)\n",
    "\n",
    "        since = time.time()\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "                for param_group in optimizer.param_groups:\n",
    "                    with open('results/Log.txt', 'a') as f:\n",
    "                        print(\"learning rate\", param_group['lr'], file=f)\n",
    "\n",
    "                model.train()  # Set model to training mode\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "\n",
    "            metrics = defaultdict(float)\n",
    "            epoch_samples = 0\n",
    "\n",
    "            for inputs, targets in dataloaders[phase]:\n",
    "                inputs = inputs.to(device)\n",
    "                targets = targets.to(device)\n",
    "\n",
    "                # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs1 = model(inputs)\n",
    "                    loss = calc_loss_dense(outputs1.float(), targets.float(), metrics)\n",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "                # track number of samples in the mini-batch\n",
    "                epoch_samples += inputs.size(0)\n",
    "            \n",
    "\n",
    "            print_metrics(metrics, epoch_samples, phase)\n",
    "            epoch_loss = metrics['loss'] / epoch_samples\n",
    "            # deep copy the model\n",
    "            if phase == 'val' and epoch_loss < best_loss:\n",
    "                print(\"saving best model\")\n",
    "                best_loss = epoch_loss\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        time_elapsed = time.time() - since\n",
    "        with open('results/Log.txt', 'a') as f:\n",
    "            print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60), file=f)\n",
    "            now = datetime.now()\n",
    "            print(\"now =\", now, file=f)\n",
    "    with open('results/Log.txt', 'a') as f:\n",
    "        print('Best val loss: {:4f}'.format(best_loss), file=f)\n",
    "   \n",
    "    # Return the best model weights according to the validation loss\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "outstanding-april",
   "metadata": {},
   "source": [
    "### Execute training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compliant-expense",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_ft = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=30, gamma=0.1)\n",
    "model = train_model(model, optimizer_ft, exp_lr_scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "unnecessary-moses",
   "metadata": {},
   "source": [
    "### Save the model with the best validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exceptional-expense",
   "metadata": {},
   "outputs": [],
   "source": [
    "stringer = 'results/' + simName + 'BestModel.pt'\n",
    "torch.save(model.state_dict(), stringer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "unique-translation",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intellectual-basement",
   "metadata": {},
   "source": [
    "### Load the model with the best validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "electronic-hunter",
   "metadata": {},
   "outputs": [],
   "source": [
    "stringer = 'results/' + simName + 'BestModel.pt'\n",
    "model.load_state_dict(torch.load(stringer))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "respiratory-curve",
   "metadata": {},
   "source": [
    "## Log file to save the test loss "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "civilian-individual",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results/' + simName + 'LogTest.txt', 'w') as f:\n",
    "    print('Test Accuracy', file=f)\n",
    "    print('-' * 20, file=f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "engaged-scanning",
   "metadata": {},
   "source": [
    "## Test Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "conceptual-requirement",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import time\n",
    "from datetime import datetime\n",
    "import copy\n",
    "from collections import defaultdict\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "\n",
    "batch_sizeTest = 1\n",
    "\n",
    "\n",
    "def my_loss(output, target):\n",
    "    loss = torch.sum((output - target)**2,1)\n",
    "    loss = torch.sqrt(loss)\n",
    "    loss = torch.mean(loss)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def calc_loss_test(pred, target, metrics):\n",
    "    loss = my_loss(pred, target)# *256*256\n",
    "    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def print_metrics(metrics, epoch_samples, phase):\n",
    "    outputs1 = []\n",
    "    for k in metrics.keys():\n",
    "        outputs1.append(\"{}: {:4f}\".format(k, metrics[k] / epoch_samples))\n",
    "    with open('results/' + simName + 'LogTest.txt', 'a') as f:\n",
    "        print(\"{}: {}\".format(phase, \", \".join(outputs1)), file=f)\n",
    "        \n",
    "def test_loss(model):\n",
    "    since = time.time()\n",
    "    model.eval()   # Set model to evaluate mode\n",
    "    metrics = defaultdict(float)\n",
    "    epoch_samples = 0\n",
    "            \n",
    "    for inputs, targets in DataLoader(Loc_test, batch_size=batch_sizeTest, shuffle=True, num_workers=8):\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "        with torch.set_grad_enabled(False):\n",
    "            outputs1 = model(inputs)\n",
    "            loss = calc_loss_test(outputs1.float(), targets.float(), metrics)\n",
    "            epoch_samples += inputs.size(0)\n",
    "     \n",
    "    print_metrics(metrics, epoch_samples, phase='test')\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    with open('results' + simName + '/LogTest.txt', 'a') as f:\n",
    "        print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60), file=f)\n",
    "        now = datetime.now()\n",
    "        print(\"now =\", now, file=f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fatal-stanley",
   "metadata": {},
   "source": [
    "### Execute test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occupied-evening",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efficient-industry",
   "metadata": {},
   "source": [
    "# Generate results as images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sitting-colleague",
   "metadata": {},
   "source": [
    "### Load the model with the best validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "desperate-extraction",
   "metadata": {},
   "outputs": [],
   "source": [
    "stringer = 'results/' + simName + 'BestModel.pt'\n",
    "model.load_state_dict(torch.load(stringer))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fatal-motor",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lib import loader, modules, modulesHeatMapOut"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "critical-terry",
   "metadata": {},
   "source": [
    "### Copy parameters to NN with heatmap output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "static-reply",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = model.state_dict()\n",
    "inp = 16\n",
    "modelMapOut = modulesHeatMapOut.LocUNet(inputs=inp)\n",
    "modelMapOut_dict = modelMapOut.state_dict()\n",
    "\n",
    "# Taken from https://discuss.pytorch.org/t/how-to-load-part-of-pre-trained-model/1113\n",
    "#Now copy and print        \n",
    "pretrained_dict = model_dict\n",
    "model2_dict = modelMapOut_dict\n",
    "# 1. filter out unnecessary keys\n",
    "pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model2_dict}\n",
    "# 2. overwrite entries in the existing state dict\n",
    "model2_dict.update(pretrained_dict)\n",
    "# 3. load the new state dict\n",
    "modelMapOut.load_state_dict(model2_dict)\n",
    "\n",
    "modelMapOut.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acting-greece",
   "metadata": {},
   "source": [
    "### Create directory for the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "incorrect-pregnancy",
   "metadata": {},
   "outputs": [],
   "source": [
    "try: \n",
    "    os.mkdir('results/img') \n",
    "except OSError as error: \n",
    "    print(error)\n",
    "    \n",
    "try: \n",
    "    os.mkdir('results/img/' + simName) \n",
    "except OSError as error: \n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affecting-single",
   "metadata": {},
   "source": [
    "### Save the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "serial-bumper",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from torch import linalg as LA\n",
    "maps_inds=np.arange(0,99,1,dtype=np.int16)\n",
    "#Standard determenistic \"random\" shuffle of the maps:\n",
    "np.random.seed(42)\n",
    "np.random.shuffle(maps_inds)\n",
    "\n",
    "\n",
    "noTrials = 50 # No of Rx considered for each map with one fixed BS deployment, max. 200\n",
    "\n",
    "from PIL import Image\n",
    "from numpy.linalg import norm\n",
    "for mapp in range(15):\n",
    "    name00=str(mapp)\n",
    "    name0=str(maps_inds[84+mapp])\n",
    "    Loc_test2 = loader.locDL(maps_inds=maps_inds, phase=\"custom\", dir_dataset=\"dataset/\",simulation=simName,\n",
    "                                           cityMap=\"true\",carsMap=\"false\",TxMaps=\"true\",\n",
    "                                           ind1=84+mapp,ind2=84+mapp\n",
    "                                           )\n",
    "    \n",
    "    ii=0\n",
    "    for inputs, target in DataLoader(Loc_test2, batch_size=1, shuffle=False, num_workers=0):\n",
    "     \n",
    "        if ii>noTrials-1:\n",
    "            break\n",
    "            \n",
    "        mapEst, est = modelMapOut(inputs.cuda())\n",
    "        \n",
    " \n",
    "        inputss=inputs.detach().cpu().numpy()\n",
    "        targets=target.detach().cpu().numpy()\n",
    "        \n",
    "        \n",
    "        mapEsts = mapEst.detach().cpu().numpy()\n",
    "        ests=est.detach().cpu().numpy()\n",
    "\n",
    "        mapEsts = mapEsts - np.min(mapEsts[0][0])\n",
    "        mapEsts = mapEsts/np.max(mapEsts[0][0])\n",
    "        mapEsts = 255*mapEsts\n",
    "\n",
    "\n",
    "        builds=inputss[0][15]\n",
    "        indB=builds!=0\n",
    "        im=np.zeros([256,256,3])\n",
    "\n",
    "        im[:,:,0] = (mapEsts[0][0])\n",
    "        im[:,:,1] = (mapEsts[0][0])\n",
    "        im[:,:,2] = (mapEsts[0][0])\n",
    "\n",
    "        name=str(ii)\n",
    "\n",
    "        im[indB,2]=200\n",
    "        im[indB,0]=0\n",
    "        im[indB,1]=0\n",
    "\n",
    "        im = Image.fromarray(im.astype(np.uint8)) \n",
    "\n",
    "        file_name=\"results/img/\" + simName + \"/\" + name00 + \"_\" + name0 + \"_\" + name + \".png\" \n",
    "        \n",
    "        im.save(file_name)\n",
    "        \n",
    "        \n",
    "        img = cv2.imread(file_name)\n",
    "        \n",
    "        aa = np.float(255)\n",
    "\n",
    "        cv2.drawMarker(img, (targets[0][1],targets[0][0]), color=(0,255,0), markerType=cv2.MARKER_CROSS, thickness=2)\n",
    "        cv2.drawMarker(img, (ests[0][1],ests[0][0]), color=(0,255,255), markerType=cv2.MARKER_SQUARE, thickness=2)\n",
    "        for tx in range(5):\n",
    "            imTx = inputss[0][10+tx]\n",
    "            ind = np.unravel_index(np.argmax(imTx, axis=None), imTx.shape)\n",
    "            cv2.drawMarker(img, (ind[1], ind[0]), color=(0,0,255), markerType=cv2.MARKER_DIAMOND, thickness=2)\n",
    "              \n",
    "        cv2.imwrite(file_name,img)\n",
    "        ii=ii+1\n"
   ]
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}