Montecarlo_Approximation_v10_Extended_Heightfield.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Montecarlo Approximation of Electron-Matter-Interaction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sys.version 3.11.0 (main, Oct 24 2022, 18:26:48) [MSC v.1933 64 bit (AMD64)]\n",
      "cuda device name(0) NVIDIA GeForce RTX 2080\n",
      "torch.__version__ 2.0.1+cu118\n",
      "fastai.__version__ 2.7.12\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "from fastai.vision.all import * \n",
    "import fastai\n",
    "import torch\n",
    "print(\"sys.version\", sys.version)\n",
    "print(\"cuda device name(0)\", torch.cuda.get_device_name(0))\n",
    "print(\"torch.__version__\", torch.__version__)\n",
    "print(\"fastai.__version__\", fastai.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\mnt\\aetna-cluster-workspace\\data\n",
      "[]\n",
      "[]\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "from os import listdir\n",
    "input_path = Path('/mnt/aetna-cluster-workspace/data/')\n",
    "\n",
    "def get_items(input_path):\n",
    "    file_names = listdir(input_path)\n",
    "    print(input_path)\n",
    "    print(file_names)\n",
    "    file_names = [filename for filename in file_names if \"_bse.tif\" in str( filename ) ]\n",
    "    return file_names\n",
    "    \n",
    "file_names = get_items(input_path)    \n",
    "print(file_names[0:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datensatz laden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "list index out of range",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 37\u001b[0m\n\u001b[0;32m     34\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(y)\n\u001b[0;32m     36\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m5\u001b[39m):\n\u001b[1;32m---> 37\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[43mfile_names\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/\u001b[39m\u001b[38;5;124m\"\u001b[39m, get_normal(file_names[i]), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/\u001b[39m\u001b[38;5;124m\"\u001b[39m, get_bse(file_names[i]))\n",
      "\u001b[1;31mIndexError\u001b[0m: list index out of range"
     ]
    }
   ],
   "source": [
    "def get_hf_0(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_entry_hf_0.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_hf_1(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_exit_hf_0.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_hf_2(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_entry_hf_1.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_hf_3(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_exit_hf_1.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_normal(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_normal.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_bse(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_bse.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "def get_se(filename):\n",
    "    filename_parts = str(filename.stem).split(\"_\")\n",
    "    y = filename.parent / Path( \"_\".join(filename_parts[:-1]) + \"_se.tif\" )\n",
    "    return str(y)\n",
    "\n",
    "for i in range(5):\n",
    "    print(file_names[i], \"/\", get_normal(file_names[i]), \"/\", get_bse(file_names[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "item_transforms  = [RandomCrop((512,512)),DihedralItem]\n",
    "\n",
    "datablocks = DataBlock(blocks=(ImageBlock(cls=PILImageBW), ImageBlock(cls=PILImageBW), ImageBlock(cls=PILImageBW), ImageBlock(cls=PILImageBW), ImageBlock, ImageBlock(cls=PILImageBW), ImageBlock(cls=PILImageBW)),\n",
    "                       n_inp=5,\n",
    "                       get_items=get_items,\n",
    "                       getters=[get_hf_0, get_hf_1, get_hf_2, get_hf_3, get_normal, get_bse, get_se],\n",
    "                       splitter=RandomSplitter(valid_pct=0.2, seed=42),\n",
    "                       item_tfms=item_transforms)\n",
    "\n",
    "data_loader = datablocks.dataloaders(input_path, bs=1, num_workers=0 )\n",
    "\n",
    "@typedispatch\n",
    "def show_batch(x:tuple, y, samples, ctxs=None, max_n=6, nrows=None, ncols=2, figsize=None, **kwargs):    \n",
    "    hf0,hf1,hf2,hf3,normal = x\n",
    "    \n",
    "    show_output = False\n",
    "    bse,se = y\n",
    "    batch_size = hf0.shape[0]\n",
    "    \n",
    "    nrows = min(x[0].shape[0], max_n)\n",
    "    if show_output:\n",
    "        ncols = 7\n",
    "    else:\n",
    "        ncols = 5\n",
    "    if ctxs is None: ctxs = get_grid(nrows*ncols, nrows=nrows, ncols=ncols, figsize=figsize)\n",
    "    \n",
    "    ctxi = 0\n",
    "    for i in range(batch_size):\n",
    "        hf0_image = hf0[i,:,:,:].squeeze(0)\n",
    "        show_image(hf0_image, title=\"heightmap (entry 1)\", ctx=ctxs[ctxi], cmap=\"gray\", **kwargs)\n",
    "        \n",
    "        ctxi = ctxi+1\n",
    "        \n",
    "        hf1_image = hf1[i,:,:,:].squeeze(0)\n",
    "        show_image(hf1_image, title=\"heightmap (exit 1)\", ctx=ctxs[ctxi], cmap=\"gray\", **kwargs)\n",
    "        ctxi = ctxi+1\n",
    "\n",
    "        hf2_image = hf2[i,:,:,:].squeeze(0)\n",
    "        show_image(hf2_image, title=\"heightmap (entry 2) \", ctx=ctxs[ctxi], cmap=\"gray\", **kwargs)\n",
    "        ctxi = ctxi+1\n",
    "        \n",
    "        hf3_image = hf3[i,:,:,:].squeeze(0)\n",
    "        show_image(hf3_image, title=\"heightmap (exit 2) \", ctx=ctxs[ctxi], cmap=\"gray\", **kwargs)\n",
    "        ctxi = ctxi+1        \n",
    "        \n",
    "        normal_image = normal[i,:,:,:].squeeze(0)      \n",
    "        show_image(normal_image, title=\"normalmap\", ctx=ctxs[ctxi], **kwargs)\n",
    "        ctxi = ctxi+1\n",
    "        \n",
    "        if show_output:\n",
    "            bse_image = bse[i,:,:,:].squeeze(0)      \n",
    "            show_image(bse_image, title=\"BSE\", ctx=ctxs[ctxi], **kwargs)\n",
    "            ctxi = ctxi+1\n",
    "\n",
    "            se_image = se[i,:,:,:].squeeze(0)      \n",
    "            show_image(se_image, title=\"SE\", ctx=ctxs[ctxi], **kwargs)\n",
    "            ctxi = ctxi+1\n",
    "    plt.savefig(\"data_representation.png\", dpi=300)  \n",
    "    \n",
    "data_loader.show_batch( )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Netzwerk Architektur"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_inner_model( dataloader, backbone, datalayout, **kwargs ):\n",
    "    img_size = dataloader.one_batch()[0].shape[-2:]\n",
    "    \n",
    "    if datalayout == \"exthf_normal\" or datalayout == \"normal_exthf\":\n",
    "        n_in=7\n",
    "    if datalayout == \"exthf_only\" or datalayout == \"hf_normal\":\n",
    "        n_in=4\n",
    "    if datalayout == \"hf_only\":\n",
    "        n_in=1\n",
    "    \n",
    "    model = create_unet_model(backbone, 2, img_size, n_in=n_in, **kwargs)\n",
    "    return model\n",
    "\n",
    "class FIBModel(torch.nn.Module):\n",
    "    def __init__(self, inner_model, datalayout = \"normal_exthf\"):\n",
    "        super().__init__()\n",
    "        self.inner_model = inner_model\n",
    "        self.datalayout = datalayout\n",
    "        \n",
    "    def forward(self, x_hf_0, x_hf_1, x_hf_2, x_hf_3, x_normal):\n",
    "\n",
    "        print(\"normal_map\", x_normal.shape, x_normal.dtype, torch.min(x_normal), \"-\", torch.max(x_normal))\n",
    "        print(\"x_hf_0    \", x_hf_0.shape, x_hf_0.dtype, torch.min(x_hf_0), \"-\", torch.max(x_hf_0))\n",
    "        print(\"x_hf_1    \", x_hf_1.shape, x_hf_1.dtype, torch.min(x_hf_1), \"-\", torch.max(x_hf_1))\n",
    "        print(\"x_hf_2    \", x_hf_2.shape, x_hf_2.dtype, torch.min(x_hf_2), \"-\", torch.max(x_hf_2))\n",
    "        print(\"x_hf_3    \", x_hf_3.shape, x_hf_3.dtype, torch.min(x_hf_3), \"-\", torch.max(x_hf_3))\n",
    "        \n",
    "        if self.datalayout == \"normal_exthf\":\n",
    "            x = torch.cat( (x_normal, x_hf_0, x_hf_1, x_hf_2, x_hf_3 ), dim=1 )\n",
    "        elif self.datalayout == \"exthf_normal\":\n",
    "            x = torch.cat( (x_hf_0, x_hf_1, x_hf_2, x_hf_3, x_normal ), dim=1 )\n",
    "        elif self.datalayout == \"exthf_only\":\n",
    "            x = torch.cat( (x_hf_0, x_hf_1, x_hf_2, x_hf_3 ), dim=1 )\n",
    "        elif self.datalayout == \"hf_normal\":\n",
    "            x = torch.cat( (x_hf_0, x_normal ), dim=1 )\n",
    "        elif self.datalayout == \"hf_only\":\n",
    "            x = x_hf_0\n",
    "        output_of_inner_model = self.inner_model(x)\n",
    "        output = torch.split(output_of_inner_model, 1, dim=1)\n",
    "        # output = output_of_inner_model\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Determine Learning Rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=0.001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.util import random_noise\n",
    "\n",
    "def add_noise( img, var ):\n",
    "    return torch.tensor( random_noise(img, mode='gaussian', mean=0, var=var, clip=True) )\n",
    "\n",
    "def estimate_snr( img ):\n",
    "    return img.mean() / img.std();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CombinedLoss():\n",
    "    def __init__(self, losses, weights=None, reduction='mean', axis=-1):\n",
    "        self.losses = losses\n",
    "        self._reduction = reduction\n",
    "        self.axis = axis\n",
    "        self.weights = weights\n",
    "        if weights is None:\n",
    "            self.weights = []\n",
    "            for _ in losses:\n",
    "                self.weights.append(1.0)\n",
    "        \n",
    "    def __call__(self, out, *yb):\n",
    "        total_loss = 0.0\n",
    "        # out = torch.split(out, 1, dim=1)\n",
    "        for i,loss_fct in enumerate( self.losses ):\n",
    "            loss_fct.reduction = 'none'\n",
    "            total_loss += loss_fct(out[i], yb[i]) * self.weights[i]\n",
    "        if self.reduction == \"mean\":\n",
    "            total_loss = total_loss.mean()\n",
    "        elif self.reduction == \"sum\":\n",
    "            total_loss = total_loss.sum()\n",
    "        return total_loss\n",
    "    \n",
    "    @property\n",
    "    def reduction(self) -> str:\n",
    "        return self._reduction    \n",
    "    \n",
    "    @reduction.setter\n",
    "    def reduction(self, reduction:str):\n",
    "        self._reduction = reduction  \n",
    "    \n",
    "    def decodes(self, x:Tensor) -> Tensor:    \n",
    "        return x\n",
    "        # return x.argmax(dim=self.axis)\n",
    "    \n",
    "    def activation(self, x:Tensor) -> Tensor:                 \n",
    "        activation = torch.zeros(x[0].shape)\n",
    "        for xi in list(x):\n",
    "            activation += F.softmax(xi, dim=self.axis)    \n",
    "        return activation  \n",
    "    \n",
    "l1          = CombinedLoss([ L1LossFlat(), L1LossFlat()] )    \n",
    "l1_weighted = CombinedLoss([ L1LossFlat(), L1LossFlat()], [1.0, 1.691]  )\n",
    "l2          = CombinedLoss([ MSELossFlat(), MSELossFlat()])\n",
    "l2_weighted = CombinedLoss([ MSELossFlat(), MSELossFlat()], [1.0, 1.691] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai import metrics\n",
    "\n",
    "def total_mse(inp, *targ):\n",
    "    total = 0.0\n",
    "    for i in range(len(inp)):\n",
    "        total += mse(inp[i], targ[i])\n",
    "    return total\n",
    "\n",
    "def total_l1(inp, *targ):\n",
    "    total = 0.0\n",
    "    for i in range(len(inp)):\n",
    "        total += mae(inp[i], targ[i])\n",
    "    return total"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mp\n",
    "cpts = [0.0, 127.0/255.0, 1.0]\n",
    "colors = [(cpts[0], (0, 0, 1.0)), (cpts[1], (0, 1.0, 0.0)), (cpts[2], (1.0, 0, 0))]\n",
    "cmap_name = 'my_list'\n",
    "colormap = mp.colors.LinearSegmentedColormap.from_list(cmap_name, colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filler(depth = 0):\n",
    "    result = \"\"\n",
    "    for i in range(depth):\n",
    "        result = result + \"  \"\n",
    "    return result\n",
    "\n",
    "def print_tuple(x, depth = 0):\n",
    "    if isinstance(x,fastai.torch_core.TensorBase):\n",
    "        print(filler(depth), type(x), x.shape)\n",
    "    elif isinstance(x,torch.Tensor):\n",
    "        print(filler(depth), type(x), x.shape)\n",
    "    elif type(x) is tuple:\n",
    "        print(filler(depth), type(x), len(x) )\n",
    "        for child_x in list(x):\n",
    "            print_tuple( child_x, depth+1 )\n",
    "    else:\n",
    "        print(filler(depth), type(x))\n",
    "\n",
    "class ConstantFunc():\n",
    "    \"Returns a function that returns `o`\"\n",
    "    def __init__(self, o): self.o = o\n",
    "    def __call__(self, *args, **kwargs): return self.o        \n",
    "        \n",
    "class PredictionsFromTupleCallback(Callback):    \n",
    "    def before_validate(self):\n",
    "        self.preds = []\n",
    "        self.targets = []\n",
    "            \n",
    "    def after_pred(self, **kwargs:Any)->None:\n",
    "        se,bse = self.pred\n",
    "        se = to_detach(se)\n",
    "        bse = to_detach(bse)\n",
    "        self.preds.append((se,bse))\n",
    "        self.targets.append(self.yb)       \n",
    "\n",
    "def create_loss_image( learner, output_filename, top=True ):\n",
    "    n_images = 6\n",
    "    \n",
    "    interpretation = Interpretation.from_learner( learner )    \n",
    "    values,indices = interpretation.top_losses(k=n_images, largest=top)\n",
    "\n",
    "    metrics_string = \"\"\n",
    "    for metric in learner.metrics:\n",
    "        metrics_string = metrics_string + metric.name + \": \" + \"{:.3f}\".format(metric.value.item() ) + \"\\n\"\n",
    "    \n",
    "    tmp_data_loader = learner.dls[1].new( get_idxs = ConstantFunc( indices ), bs=1 )\n",
    "    cb = PredictionsFromTupleCallback()\n",
    "    ctx_mgrs = learner.validation_context(cbs=[cb])\n",
    "    with ContextManagers(ctx_mgrs):\n",
    "        learner._do_epoch_validate(dl=tmp_data_loader)\n",
    "\n",
    "    all_predictions = cb.preds\n",
    "    all_targets     = cb.targets\n",
    "\n",
    "    figure      = plt.figure( constrained_layout=True, figsize=(16,12*n_images) )\n",
    "    figure.suptitle(metrics_string, fontsize=16 )\n",
    "    \n",
    "    sub_figures = figure.subfigures(nrows=n_images, ncols=1)\n",
    "    \n",
    "    # fig, axs = plt.subplots(2*4, 3, figsize=(16,48))\n",
    "\n",
    "    for i,(idx,loss_value) in enumerate(zip(indices,values)):\n",
    "        filename = str( learner.dls.valid_ds.items[idx] )\n",
    "        \n",
    "        sub_figures[i].suptitle( filename + \"\\nLoss:{:.3f}\".format(loss_value) )\n",
    "        axs = sub_figures[i].subplots(nrows=2, ncols=3)        \n",
    "        \n",
    "        # hf_preds,bse_preds = preds\n",
    "        se_pred,bse_pred = all_predictions[i]\n",
    "        se_pred  = torch.squeeze( se_pred, 0 )\n",
    "        bse_pred = torch.squeeze( bse_pred, 0 )\n",
    "\n",
    "        se_target,bse_target = all_targets[i]\n",
    "        se_target  = torch.squeeze(  se_target, 0 )\n",
    "        bse_target = torch.squeeze( bse_target, 0 )\n",
    "\n",
    "        se_resid  = abs(se_pred.cpu() - se_target.cpu())\n",
    "        bse_resid = abs(bse_pred.cpu() - bse_target.cpu())\n",
    "\n",
    "        show_image( ax=axs[0,0], im=se_pred,        title=\"SE Prediction \", vmin=0, vmax=1,  cmap=\"gray\")\n",
    "        show_image( ax=axs[1,0], im=bse_pred,       title=\"BSE Prediction\", vmin=0, vmax=1,  cmap=\"gray\")\n",
    "\n",
    "        show_image( ax=axs[0,1], im=se_target,      title=\"SE Target\", vmin=0, vmax=1,       cmap=\"gray\")\n",
    "        show_image( ax=axs[1,1], im=bse_target,     title=\"BSE Target\", vmin=0, vmax=1,      cmap=\"gray\")\n",
    "\n",
    "        show_image( ax=axs[0,2], im=se_resid,       title=\"SE Residual\", vmin=0, vmax=0.20,  cmap=\"coolwarm\")\n",
    "        show_image( ax=axs[1,2], im=bse_resid,      title=\"BSE Residual\", vmin=0, vmax=0.20, cmap=\"coolwarm\")\n",
    "\n",
    "    plt.savefig(output_filename, dpi=150)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.metrics import structural_similarity as ssim\n",
    "from skimage.metrics import peak_signal_noise_ratio as psnr\n",
    "import time\n",
    "\n",
    "def get_predictions( learner ):\n",
    "    callback = PredictionsFromTupleCallback()\n",
    "    ctx_mgrs = learner.validation_context(cbs=[callback])\n",
    "    with ContextManagers(ctx_mgrs):\n",
    "        learner._do_epoch_validate(dl=learner.dls.valid)\n",
    "    return callback.preds,callback.targets\n",
    "\n",
    "def evaluate_predictions( learner ):\n",
    "    print(learner)\n",
    "    \n",
    "    print(\"performing predictions\")\n",
    "    start = time.perf_counter()\n",
    "    predictions,targets = get_predictions( learner )\n",
    "    stop = time.perf_counter()\n",
    "    print(\"done performing predictions of \", len(predictions), f\"items in {stop-start:0.4f}\")\n",
    "\n",
    "    n_images = 0\n",
    "    \n",
    "    bse_std = 0.0\n",
    "    se_std  = 0.0\n",
    "       \n",
    "    ssim_bse = 0.0\n",
    "    ssim_se  = 0.0\n",
    "    \n",
    "    psnr_bse = 0.0\n",
    "    psnr_se  = 0.0\n",
    "        \n",
    "    for (bse_pred,se_pred),(bse_target,se_target) in zip(predictions,targets) :\n",
    "        bse_pred   = bse_pred.cpu().to(dtype=torch.float32)\n",
    "        se_pred    = se_pred.cpu().to(dtype=torch.float32)\n",
    "        bse_target = bse_target.cpu().to(dtype=torch.float32)\n",
    "        se_target  = se_target.cpu().to(dtype=torch.float32)\n",
    "        \n",
    "        bse_pred   = bse_pred * 255.0\n",
    "        se_pred    = se_pred * 255.0\n",
    "        bse_target = bse_target * 255.0\n",
    "        se_target  = se_target * 255.0\n",
    "        \n",
    "        bse_residual = bse_pred - bse_target\n",
    "        se_residual  = se_pred - se_target\n",
    "       \n",
    "        for i in range(bse_residual.shape[0]):\n",
    "            n_images = n_images + 1\n",
    "            bse_std = bse_std + torch.std(bse_residual[i,0,:,:])            \n",
    "            se_std = se_std + torch.std(se_residual[i,0,:,:])\n",
    "            \n",
    "            bse_a = bse_pred[i,0,:,:].squeeze(0).squeeze(0).cpu().detach().numpy()\n",
    "            bse_b = bse_target[i,0,:,:].squeeze(0).squeeze(0).cpu().detach().numpy()            \n",
    "            ssim_bse = ssim_bse + ssim(bse_a, bse_b, data_range=255)\n",
    "            psnr_bse = psnr_bse + psnr(bse_a, bse_b, data_range=255)\n",
    "\n",
    "            se_a = se_pred[i,0,:,:].squeeze(0).squeeze(0).cpu().detach().numpy()\n",
    "            se_b = se_target[i,0,:,:].squeeze(0).squeeze(0).cpu().detach().numpy()            \n",
    "            ssim_se = ssim_se + ssim(se_a, se_b, data_range=255)\n",
    "            psnr_se = psnr_se + psnr(se_a, se_b, data_range=255)\n",
    "            \n",
    "    print( \"bse std\",  \"{:.2f}\".format(bse_std / float(n_images) ) )\n",
    "    print( \"se  std \", \"{:.2f}\".format(se_std / float(n_images) ) )\n",
    "    \n",
    "    print( \"bse ssim\",  \"{:.2f}\".format(ssim_bse / float(n_images) ) )\n",
    "    print( \"se  ssim \", \"{:.2f}\".format(ssim_se / float(n_images) ) )\n",
    "    \n",
    "    print( \"bse psnr\",  \"{:.2f}\".format(psnr_bse / float(n_images) ) )\n",
    "    print( \"se  psnr \", \"{:.2f}\".format(psnr_se / float(n_images) ) )\n",
    "            \n",
    "# evaluate_predictions( learner )            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_losses =  [ (l1, \"l1\"), (l2, \"l2\"), (l1_weighted, \"l1_weighted\"), (l2_weighted, \"l2_weighted\") ]\n",
    "all_backbones = [ (resnet152, \"resnet152\"), (resnet101, \"resnet101\"), (resnet50,\"resnet50\"), (resnet34,\"resnet34\") ]\n",
    "all_datalayout = [\"exthf_normal\", \"normal_exthf\", \"exthf_only\", \"hf_normal\", \"hf_only\" ]\n",
    "\n",
    "experiments = []\n",
    "experiments.append( ( (l2, \"l2\"), (resnet152, \"resnet152\"), \"exthf_normal\") )\n",
    "\n",
    "# for loss in all_losses:\n",
    "#    experiments.append( (loss, (resnet152, \"resnet152\"), \"exthf_normal\") )#\n",
    "#for backbone in all_backbones:\n",
    "#    experiments.append(  ((l1, \"l1\"), backbone, \"exthf_normal\") )\n",
    "#for layout in all_datalayout:\n",
    "#    experiments.append(  ((l1, \"l1\"), (resnet152, \"resnet152\"), layout ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train    = True\n",
    "evaluate = True\n",
    "test     = False\n",
    "\n",
    "n_epochs = [19,30,50]\n",
    "\n",
    "model_dir = \"/mnt/aetna-cluster-workspace/models/\"\n",
    "    \n",
    "for (loss_func, lossname),(backbone, backbone_name),datalayout in experiments:\n",
    "    model = FIBModel( create_inner_model( data_loader, backbone, datalayout ), datalayout )\n",
    "    learner = Learner( data_loader, model, model_dir=model_dir, loss_func=loss_func, metrics=[total_mse, total_l1] )\n",
    "\n",
    "    if train:\n",
    "        learner.fit( 1, lr=learning_rate )\n",
    "        epochs = 1\n",
    "\n",
    "        for step in n_epochs:\n",
    "            epochs = epochs + step\n",
    "            weight_name = datalayout + \"_\" + backbone_name + \"_\" + lossname + \"_\" + str(epochs)\n",
    "            print(weight_name)\n",
    "            learner.fit( step, lr=learning_rate, cbs=[ShowGraphCallback(), CSVLogger(filename=weight_name+\".csv\")] )    \n",
    "            learner.save( weight_name )\n",
    "    else:\n",
    "        epochs = 100\n",
    "        weight_name = datalayout + \"_\" + backbone_name + \"_\" + lossname + \"_\" + str(epochs)\n",
    "        print(\"loading weights\", weight_name)\n",
    "        learner.load(weight_name)\n",
    "\n",
    "    if evaluate:\n",
    "        evaluate_predictions( learner )\n",
    "        create_loss_image( learner, weight_name + \"_worst6.png\", top=True)\n",
    "        create_loss_image( learner, weight_name + \"_bestp6.png\", top=False)\n",
    "        \n",
    "    if test:\n",
    "        test_input_path = Path('./test_data')\n",
    "        test_files = get_items( test_input_path )\n",
    "        test_dataloader = learner.dls.test_dl( test_files )\n",
    "\n",
    "        cb = PredictionsFromTupleCallback()\n",
    "        ctx_mgrs = learner.validation_context(cbs=[cb])\n",
    "        with ContextManagers(ctx_mgrs):\n",
    "            learner._do_epoch_validate(dl=test_dataloader)\n",
    "        all_preds = cb.preds\n",
    "\n",
    "        # all_preds = learner.get_preds( dl=test_dataloader, with_decoded=False, with_input=True )            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_rows = 8\n",
    "\n",
    "figure   = plt.figure  ( constrained_layout=True )\n",
    "fig, axs = plt.subplots( nrows=n_rows, ncols=2, figsize=(16,8*n_rows) )        \n",
    "\n",
    "for i,preds in enumerate(all_preds[0:n_rows]):\n",
    "    se,bse = preds\n",
    "    show_image( ax=axs[i,0], im=se[0,:,:],  title=\"SE\", cmap=\"gray\")    \n",
    "    show_image( ax=axs[i,1], im=bse[0,:,:], title=\"BSE\",  cmap=\"gray\")"
   ]
  }
 ],
 "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}