plot_examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import models.unet_precip_regression_lightning as unet_regr\n",
    "from utils import dataset_precip\n",
    "\n",
    "dataset = dataset_precip.precipitation_maps_oversampled_h5(\n",
    "            in_file=\"data/precipitation/train_test_2016-2019_input-length_12_img-ahead_6_rain-threshold_50.h5\",\n",
    "            num_input_images=12,\n",
    "            num_output_images=6, train=False)\n",
    "# 222, 444, 777, 1337 is fine\n",
    "x, y_true = dataset[777]\n",
    "# add batch dimension\n",
    "x = torch.tensor(x).unsqueeze(0)\n",
    "y_true = torch.tensor(y_true).unsqueeze(0)\n",
    "print(x.shape, y_true.shape)\n",
    "\n",
    "models_dir = \"lightning/precip_regression/comparison\"\n",
    "precip_models = [m for m in os.listdir(models_dir) if \".ckpt\" in m]\n",
    "print(precip_models)\n",
    "\n",
    "loss_func = nn.MSELoss(reduction=\"sum\")\n",
    "print(\"Persistence\", loss_func(x.squeeze()[-1]*47.83, y_true.squeeze()*47.83).item())\n",
    "\n",
    "plot_images = dict()\n",
    "for i, model_file in enumerate(precip_models):\n",
    "    # model_name = \"\".join(model_file.split(\"_\")[0])\n",
    "    if \"UNet_Attention\" in model_file:\n",
    "        model_name = \"UNet with CBAM\"\n",
    "        model = unet_regr.UNet_Attention\n",
    "    elif \"UNetDS_Attention_4CBAMs\" in model_file:\n",
    "        model_name = \"UNetDS Attention 4CBAMs\"\n",
    "        model = unet_regr.UNetDS_Attention_4CBAMs\n",
    "    elif \"UNetDS_Attention\" in model_file:\n",
    "        model_name = \"SmaAt-UNet\"\n",
    "        model = unet_regr.UNetDS_Attention\n",
    "    elif \"UNetDS\" in model_file:\n",
    "        model_name = \"UNet with DSC\"\n",
    "        model = unet_regr.UNetDS\n",
    "    elif \"UNet\" in model_file:\n",
    "        model_name = \"UNet\"\n",
    "        model = unet_regr.UNet\n",
    "    else:\n",
    "        raise NotImplementedError(f\"Model not found\")\n",
    "    model = model.load_from_checkpoint(f\"{models_dir}/{model_file}\")\n",
    "    # loaded = torch.load(f\"{models_dir}/{model_file}\")\n",
    "    # model = loaded[\"model\"]\n",
    "    # model.load_state_dict(loaded[\"state_dict\"])\n",
    "    # loss_func = nn.CrossEntropyLoss(weight=torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.]))\n",
    "\n",
    "    model.to(\"cpu\")\n",
    "    model.eval()\n",
    "\n",
    "    # get prediction of model\n",
    "    with torch.no_grad():\n",
    "        y_pred = model(x)\n",
    "    y_pred = y_pred.squeeze()\n",
    "    print(model_name, loss_func(y_pred*47.83, y_true.squeeze()*47.83).item())\n",
    "\n",
    "    plot_images[model_name] = y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#### Plotting ####\n",
    "plt.figure(figsize=(10, 20))\n",
    "plt.subplot(len(precip_models)+2, 1, 1)\n",
    "plt.imshow(y_true.squeeze()*47.83)\n",
    "plt.colorbar()\n",
    "plt.title(\"Ground truth\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(len(precip_models)+2, 1, 2)\n",
    "plt.imshow(x.squeeze()[-1]*47.83, vmax=torch.max(y_true*47.83), vmin=torch.min(y_true*47.83))\n",
    "plt.colorbar()\n",
    "plt.title(\"Persistence\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "for i, (model_name, y_pred) in enumerate(plot_images.items()):\n",
    "    print(model_name)\n",
    "    plt.subplot(len(precip_models)+2, 1, i+3)\n",
    "    plt.imshow(y_pred*47.83, vmax=torch.max(y_true*47.83), vmin=torch.min(y_true*47.83))\n",
    "    plt.colorbar()\n",
    "    plt.title(model_name)\n",
    "    plt.axis(\"off\")\n",
    "# plt.colorbar()\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, axesgrid = plt.subplots(nrows=2, ncols=(len(precip_models)+2)//2, figsize=(20, 10))\n",
    "axes = axesgrid.flat\n",
    "im = axes[0].imshow(y_true.squeeze()*47.83)\n",
    "axes[0].set_title(\"Groundtruth\", fontsize=15)\n",
    "axes[0].set_axis_off()\n",
    "axes[1].imshow(x.squeeze()[-1]*47.83, vmax=torch.max(y_true*47.83), vmin=torch.min(y_true*47.83))\n",
    "axes[1].set_title(\"Persistence\", fontsize=15)\n",
    "axes[1].set_axis_off()\n",
    "\n",
    "# for ax in axes.flat[2:]:\n",
    "for i, (model_name, y_pred) in enumerate(plot_images.items()):\n",
    "    print(model_name)\n",
    "#     plt.subplot(len(precip_models)+2, 1, i+3)\n",
    "    axes[i+2].imshow(y_pred*47.83, vmax=torch.max(y_true*47.83), vmin=torch.min(y_true*47.83))\n",
    "    axes[i+2].set_title(model_name, fontsize=15)\n",
    "    axes[i+2].set_axis_off()\n",
    "#     im = ax.imshow(np.random.random((16, 16)), cmap='viridis',\n",
    "#                    vmin=0, vmax=1)\n",
    "\n",
    "fig.subplots_adjust(wspace=0.02)\n",
    "cbar = fig.colorbar(im, ax=axesgrid.ravel().tolist(), shrink=1)\n",
    "cbar.ax.set_ylabel('mm/5min', fontsize=15)\n",
    "plt.savefig(\"plots/Precip_example.eps\")\n",
    "plt.savefig(\"plots/Precip_example.png\", dpi=300)\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}