module-usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ISBISegment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.datasets import ISBISegment\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "ds = ISBISegment(\"./datasets/ISBI-2012-challenge/\", trim_seg=True)\n",
    "\n",
    "sample = ds[0]\n",
    "img, seg = sample[\"img\"], sample[\"seg\"]\n",
    "orig_img = img[\n",
    "    ...,\n",
    "    ds.subject_region[1] : ds.subject_region[3],\n",
    "    ds.subject_region[0] : ds.subject_region[2],\n",
    "]\n",
    "print(f\"Image shape: {img.shape}, dtype: {img.dtype}\")\n",
    "print(f\"Segmentation shape: {seg.shape}, dtype: {seg.dtype}\")\n",
    "print(f\"Sample count: {len(ds)}\")\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(12,4))\n",
    "ax[0].imshow(orig_img.numpy().transpose(1, 2, 0))\n",
    "ax[1].imshow(seg.numpy().transpose(1, 2, 0), cmap=\"gray\")\n",
    "ax[2].imshow(img.numpy().transpose(1, 2, 0))\n",
    "\n",
    "ax[0].set_title(\"Original image\")\n",
    "ax[1].set_title(\"Segmentation (target)\")\n",
    "ax[2].set_title(\"Extrapolated image (model input)\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ISBICellTrack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.datasets import ISBICellTrack\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "ds = ISBICellTrack(\"./datasets/PhC-C2DH-U373\")  # or \"datasets/DIC-C2DH-HeLa\"\n",
    "\n",
    "sample = ds[0]\n",
    "img, seg = sample[\"img\"], sample[\"seg\"]\n",
    "orig_img = img[\n",
    "    ...,\n",
    "    ds.subject_region[1] : ds.subject_region[3],\n",
    "    ds.subject_region[0] : ds.subject_region[2],\n",
    "]\n",
    "print(f\"Image shape: {img.shape}, dtype: {img.dtype}\")\n",
    "print(f\"Segmentation shape: {seg.shape}, dtype: {seg.dtype}\")\n",
    "print(f\"Sample count: {len(ds)}\")\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(12, 4))\n",
    "ax[0].imshow(orig_img.numpy().transpose(1, 2, 0))\n",
    "ax[1].imshow(seg.numpy().transpose(1, 2, 0), cmap=\"gray\")\n",
    "ax[2].imshow(img.numpy().transpose(1, 2, 0))\n",
    "\n",
    "ax[0].set_title(\"Original image\")\n",
    "ax[1].set_title(\"Segmentation (target)\")\n",
    "ax[2].set_title(\"Extrapolated image (model input)\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from src.learner import UNet\n",
    "\n",
    "model = UNet()\n",
    "img = torch.Tensor(8, 1, 572, 572)\n",
    "batch = {\"img\": img}\n",
    "out = model(batch)\n",
    "out[\"logits\"].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.datasets import ISBISegment\n",
    "from src.learner import UNet\n",
    "from src.loss import SegmentCrossEntropy\n",
    "\n",
    "\n",
    "sample_id = 0\n",
    "device = 1\n",
    "ds = ISBISegment(\"./datasets/ISBI-2012-challenge/\")\n",
    "model = UNet()\n",
    "model.set_devices([device])\n",
    "loss_fn = SegmentCrossEntropy(device, [92, 92, 480, 480])\n",
    "sample = ds[sample_id]\n",
    "info = model(sample)\n",
    "loss = loss_fn(info, sample)\n",
    "\n",
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from src.util.data import get_9_pt_flow, warp_image\n",
    "from src.util.visualize import flow2rgb\n",
    "from src.datasets import ISBISegment\n",
    "\n",
    "\n",
    "ds = ISBISegment(\"./datasets/ISBI-2012-challenge/\", do_aug=False)\n",
    "sample = ds[0]\n",
    "img = sample[\"img\"].numpy().transpose(1, 2, 0)\n",
    "\n",
    "flow = get_9_pt_flow((572, 572), std=10)  # also can use random_warp\n",
    "warped_img = warp_image(img, flow)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(12, 4))\n",
    "\n",
    "ax[0].imshow(img.squeeze())\n",
    "ax[1].imshow(flow2rgb(flow))\n",
    "ax[2].imshow(warped_img.squeeze())\n",
    "\n",
    "ax[0].set_title(\"Original image (extrapolated)\")\n",
    "ax[1].set_title(\"Warping flow\")\n",
    "ax[2].set_title(\"Warped image\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from mt_pipe.src.evaluators import SegmentationEvaluator\n",
    "from src.datasets import ISBICellTrack\n",
    "from src.learner import UNet\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "device = 1\n",
    "ds = ISBICellTrack(\"./datasets/DIC-C2DH-HeLa\")\n",
    "dl = DataLoader(ds, 4)\n",
    "model = UNet()\n",
    "# ckpt = torch.load(\n",
    "#     \"out/isbi-cell_track-dic_hela/run0/ckpts/best.ckpt\"\n",
    "# )  # load checkpoint if available\n",
    "# model.load_state_dict(ckpt[\"learner\"])\n",
    "model.set_devices([device])\n",
    "evaluator = SegmentationEvaluator(batch_img_key=\"img_orig\")\n",
    "# evaluator.set_out_path(\"temp\")\n",
    "results = []\n",
    "for i, batch in enumerate(dl):\n",
    "    info = model(batch)\n",
    "    res = evaluator.process_batch(batch, info)\n",
    "    results.append(res)\n",
    "    # if i >= 2:\n",
    "    #     break\n",
    "\n",
    "evaluator.output(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ie",
   "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
}