{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "from rtnls_inference import (\n",
    "    HeatmapRegressionEnsemble,\n",
    "    SegmentationEnsemble,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Segmentation of preprocessed images\n",
    "\n",
    "Here we segment images preprocessed using 0_preprocess.ipynb\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_path = Path(\"../samples/fundus\")\n",
    "\n",
    "# input folders. these are the folders where we stored the preprocessed images\n",
    "rgb_path = ds_path / \"rgb\"\n",
    "ce_path = ds_path / \"ce\"\n",
    "\n",
    "# these are the output folders for:\n",
    "av_path = ds_path / \"av\"  # artery-vein segmentations\n",
    "discs_path = ds_path / \"discs\"  # optic disc segmentations\n",
    "overlays_path = ds_path / \"overlays\"  # optional overlay visualizations\n",
    "\n",
    "device = torch.device(\"cuda:0\")  # device to use for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_paths = sorted(list(rgb_path.glob(\"*.png\")))\n",
    "ce_paths = sorted(list(ce_path.glob(\"*.png\")))\n",
    "paired_paths = list(zip(rgb_paths, ce_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paired_paths[0]  # important to make sure that the paths are paired correctly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Artery-vein segmentation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "av_ensemble = SegmentationEnsemble.from_huggingface('Eyened/vascx:artery_vein/av_july24.pt').to(device)\n",
    "\n",
    "av_ensemble.predict_preprocessed(paired_paths, dest_path=av_path, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Disc segmentation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_ensemble = SegmentationEnsemble.from_huggingface('Eyened/vascx:disc/disc_july24.pt').to(device)\n",
    "disc_ensemble.predict_preprocessed(paired_paths, dest_path=discs_path, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fovea detection\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fovea_ensemble = HeatmapRegressionEnsemble.from_huggingface('Eyened/vascx:fovea/fovea_july24.pt').to(device)\n",
    "# note: this model does not use contrast enhanced images\n",
    "df = fovea_ensemble.predict_preprocessed(paired_paths, num_workers=2)\n",
    "df.columns = [\"mean_x\", \"mean_y\"]\n",
    "df.to_csv(ds_path / \"fovea.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the retinas (optional)\n",
    "\n",
    "This will only work if you ran all the models and stored the outputs using the same folder/file names as above\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vascx.fundus.loader import RetinaLoader\n",
    "\n",
    "from rtnls_enface.utils.plotting import plot_gridfns\n",
    "\n",
    "loader = RetinaLoader.from_folder(ds_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_gridfns([ret.plot for ret in loader[:6]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Storing visualizations (optional)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not overlays_path.exists():\n",
    "    overlays_path.mkdir()\n",
    "for ret in loader:\n",
    "    fig, _ = ret.plot()\n",
    "    fig.savefig(overlays_path / f\"{ret.id}.png\", bbox_inches=\"tight\", pad_inches=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "retinalysis",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}