visualize_results.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72e163ad",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Visualize Results of a Patch-Based Model\n",
    "This notebook allows you to see the results of a patch-based model.<br>\n",
    "You can look at just sinograms, or both sinograms and reconstructions.<br>\n",
    "You can also choose which sample you want to see results from, and which slices you want to see.<br>\n",
    "The default behaviour is to randomly choose a sample number and slices.<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23c6ac2f",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "025e2a89",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import wandb\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP, DataParallel as DP\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"figure.figsize\"] = (12, 8)\n",
    "plt.rcParams[\"figure.dpi\"] = 300\n",
    "%matplotlib inline\n",
    "\n",
    "from network.patch_visualizer import PatchVisualizer\n",
    "from network.models import MaskedGAN\n",
    "from network.models.generators import PatchUNet\n",
    "from network.models.discriminators import PatchDiscriminator\n",
    "\n",
    "from utils.tomography import reconstruct\n",
    "\n",
    "# Number of threads for OpenMP. If too high, may cause error\n",
    "%env OMP_NUM_THREADS=16\n",
    "\n",
    "# Auto-reload code from disk\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb48d8c",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "Specify the parameters you want to run with here.<br>\n",
    "\n",
    "`data_dir` is the directory containing the input data.<br>\n",
    "`model_file` is the path to the model you want to visualize.<br>\n",
    "`mask_file` is the path to the binary mask indicating locations of stripes.<br>\n",
    "`sample_no` is the sample number to load data from.<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5170f4e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "i12 = Path('/dls/i12/data/2022/nt33730-1/processing/NoStripesNet')\n",
    "data_dir = i12/'data'/'wider_stripes'\n",
    "model_file = i12/'pretrained_models'/'five_sample'/'4x4'/'val'/'five_sample_4x4_100.tar'\n",
    "mask_file = i12/'stripe_masks.npz'\n",
    "\n",
    "rng = np.random.default_rng()\n",
    "cor = 1253\n",
    "sample_no = 0\n",
    "print(f\"Sample No.: {sample_no}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111dab02",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "090b6204",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    d = torch.device('cuda')\n",
    "else:\n",
    "    d = torch.device('cpu')\n",
    "\n",
    "# Load model state dict from disk\n",
    "checkpoint = torch.load(model_file, map_location=d)\n",
    "# Initialize Generator and Discriminator\n",
    "gen = DP(PatchUNet())\n",
    "gen.load_state_dict(checkpoint['gen_state_dict'])\n",
    "disc = DP(PatchDiscriminator())\n",
    "disc.load_state_dict(checkpoint['disc_state_dict'])\n",
    "# Initialize Model\n",
    "model = MaskedGAN(gen, disc, mode='test', device=d)\n",
    "\n",
    "# Initialize Visualizer\n",
    "v = PatchVisualizer(data_dir, model, sample_no=sample_no, mask_file=mask_file)\n",
    "\n",
    "# Figure Dictionary\n",
    "fig_dict = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ae40022",
   "metadata": {},
   "source": [
    "## Plot Fake Artifacts\n",
    "\n",
    "`fake_artifact_idx` is the slice index of the sinogram you want to plot.<br>\n",
    "Be sure to choose a sinogram that has **no** real-life artifacts, as otherwise you'll see black boxes where the real artifact is.<br>\n",
    "The `clean_idxs` attribute of `PatchVisualizer` is a list of all sinogram slice indexes that don't contain any real-life artifacts.<br>\n",
    "By default, `fake_artifact_idx` is randomly chosen from `clean_idxs`.<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ce10b78",
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_artifact_idx = rng.choice(v.clean_idxs)\n",
    "print(f\"Index: {fake_artifact_idx}\")\n",
    "clean = v.get_sinogram(fake_artifact_idx, 'clean')\n",
    "stripe = v.get_sinogram(fake_artifact_idx, 'stripe')\n",
    "gen_out = v.get_model_sinogram(fake_artifact_idx, 'fake')\n",
    "\n",
    "clean_r = reconstruct(clean, rot_center=cor, ncore=16)\n",
    "stripe_r = reconstruct(stripe, rot_center=cor, ncore=16)\n",
    "gen_out_r = reconstruct(gen_out, rot_center=cor, ncore=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5da307",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 3)\n",
    "fig.suptitle(\"Synthetic Stripes\", size='xx-large')\n",
    "\n",
    "axs[0, 0].set_title(f\"Clean {fake_artifact_idx}\")\n",
    "axs[0, 0].axis('off')\n",
    "axs[0, 0].imshow(clean, cmap='gray')\n",
    "\n",
    "axs[0, 1].set_title(f\"Stripe {fake_artifact_idx}\")\n",
    "axs[0, 1].axis('off')\n",
    "axs[0, 1].imshow(stripe, cmap='gray')\n",
    "\n",
    "axs[0, 2].set_title(f\"Model Output {fake_artifact_idx}\")\n",
    "axs[0, 2].axis('off')\n",
    "axs[0, 2].imshow(gen_out, cmap='gray')\n",
    "\n",
    "axs[1, 0].axis('off')\n",
    "axs[1, 0].imshow(clean_r, cmap='gray', vmin=-0.03, vmax=0.15)\n",
    "axs[1, 1].axis('off')\n",
    "axs[1, 1].imshow(stripe_r, cmap='gray', vmin=-0.03, vmax=0.15)\n",
    "axs[1, 2].axis('off')\n",
    "axs[1, 2].imshow(gen_out_r, cmap='gray', vmin=-0.03, vmax=0.15)\n",
    "\n",
    "fig_dict[f\"Sinogram {fake_artifact_idx}\"] = fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "102bef38",
   "metadata": {},
   "source": [
    "## Plot Real-life Artifacts\n",
    "`real_artifact_idx` is the slice index of the sinogram you want to plot.<br>\n",
    "Be sure to choose a sinogram that has **at least one** real-life artifact, as otherwise you won't be able to see the effect of the model.<br>\n",
    "The `stripe_idxs` attribute of `PatchVisualizer` is a list of all sinogram slice indexes that contain at least one real-life artifact.<br>\n",
    "By default, `real_artifact_idx` is randomly chosen from `stripe_idxs`.<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89339920",
   "metadata": {},
   "outputs": [],
   "source": [
    "real_artifact_idx = rng.choice(v.stripe_idxs)\n",
    "print(f\"Index: {real_artifact_idx}\")\n",
    "stripe = v.get_sinogram(real_artifact_idx, 'raw')\n",
    "gen_out = v.get_model_sinogram(real_artifact_idx, 'real')\n",
    "\n",
    "stripe_r = reconstruct(stripe, rot_center=cor, ncore=16)\n",
    "gen_out_r = reconstruct(gen_out, rot_center=cor, ncore=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fc6092",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2)\n",
    "fig.suptitle(\"Real-life Stripes\", size='xx-large')\n",
    "\n",
    "axs[0, 0].set_title(f\"Stripe {real_artifact_idx}\")\n",
    "axs[0, 0].axis('off')\n",
    "axs[0, 0].imshow(stripe, cmap='gray')\n",
    "\n",
    "axs[0, 1].set_title(f\"Model Output {real_artifact_idx}\")\n",
    "axs[0, 1].axis('off')\n",
    "axs[0, 1].imshow(gen_out, cmap='gray')\n",
    "\n",
    "axs[1, 0].axis('off')\n",
    "axs[1, 0].imshow(stripe_r, cmap='gray', vmin=-0.03, vmax=0.15)\n",
    "axs[1, 1].axis('off')\n",
    "axs[1, 1].imshow(gen_out_r, cmap='gray', vmin=-0.03, vmax=0.15)\n",
    "\n",
    "fig_dict[f\"Sinogram {real_artifact_idx}\"] = fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bef010c0",
   "metadata": {},
   "source": [
    "## Upload to Weights & Biases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f58e39e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f3924c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "api = wandb.Api()\n",
    "runs = api.runs(\"nostripesnet/NoStripesNet\")\n",
    "run_name = \"\"\n",
    "for rns in runs:\n",
    "    if rns.name == run_name:\n",
    "        wandb.init(project='NoStripesNet', entity='nostripesnet', id=rns.id, resume='must')\n",
    "        wandb.log(fig_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5c0daf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}