{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for DynamicUnet:\n\tsize mismatch for layers.12.0.weight: copying a param with shape torch.Size([5, 291, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 291, 1, 1]).\n\tsize mismatch for layers.12.0.bias: copying a param with shape torch.Size([5]) from checkpoint, the shape in current model is torch.Size([3]).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-5001b501bd01>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"cuda\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 35\u001b[0;31m \u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'ecadBlur3Watershed_101'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     36\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/fastai/lib/python3.8/site-packages/fastai/learner.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(self, file, device, **kwargs)\u001b[0m\n\u001b[1;32m    357\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcreate_opt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    358\u001b[0m     \u001b[0mfile\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjoin_path_file\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel_dir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mext\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'.pth'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 359\u001b[0;31m     \u001b[0mload_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    360\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    361\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/fastai/lib/python3.8/site-packages/fastai/learner.py\u001b[0m in \u001b[0;36mload_model\u001b[0;34m(file, model, opt, with_opt, device, strict)\u001b[0m\n\u001b[1;32m     49\u001b[0m     \u001b[0mhasopt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m'model'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'opt'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m     \u001b[0mmodel_state\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mstate\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'model'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mhasopt\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mstate\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 51\u001b[0;31m     \u001b[0mget_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstrict\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstrict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     52\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mhasopt\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mwith_opt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'opt'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/fastai/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[0;34m(self, state_dict, strict)\u001b[0m\n\u001b[1;32m   1049\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1050\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1051\u001b[0;31m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[0m\u001b[1;32m   1052\u001b[0m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[1;32m   1053\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for DynamicUnet:\n\tsize mismatch for layers.12.0.weight: copying a param with shape torch.Size([5, 291, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 291, 1, 1]).\n\tsize mismatch for layers.12.0.bias: copying a param with shape torch.Size([5]) from checkpoint, the shape in current model is torch.Size([3])."
     ]
    }
   ],
   "source": [
    "from fastai.vision.all import *\n",
    "import cv2\n",
    "import skimage as sm\n",
    "import skimage.io\n",
    "import tifffile\n",
    "path = Path('/notebooks/EcadWatershedTrain')\n",
    "# Training\n",
    "\n",
    "valid_fnames = (path/'valid.txt').read_text().split('\\n')\n",
    "path_im = path/'dat/ecadBlur3_101/train_images'\n",
    "path_lbl = path/'dat/ecadBlur3_101/train_masks'\n",
    "codes = np.loadtxt(path/'codes.txt', dtype=str)\n",
    "# codes = np.loadtxt(path/'codes.txt', dtype=str)\n",
    "fnames = get_image_files(path_im)\n",
    "lbl_names = get_image_files(path_lbl)\n",
    "get_msk = lambda o: path/f'dat/ecadBlur3_101/train_masks'/f'mask{o.stem.split(\"g\")[1]}{o.suffix}'\n",
    "\n",
    "camvid = DataBlock(blocks=(ImageBlock, MaskBlock(codes)),\n",
    "                   get_items=get_image_files,\n",
    "                   splitter=FileSplitter(path/'valid.txt'),\n",
    "                   get_y=get_msk,\n",
    "                   batch_tfms=[*aug_transforms(pad_mode='zeros'), Normalize.from_stats(*imagenet_stats)])\n",
    "\n",
    "dls = camvid.dataloaders(path/f'dat/ecadBlur3_101/train_images', bs=4)\n",
    "dls.vocab = codes\n",
    "name2id = {v:k for k,v in enumerate(codes)}\n",
    "void_code = name2id['Void']\n",
    "def acc_camvid(inp, targ):\n",
    "  targ = targ.squeeze(1)\n",
    "  mask = targ != void_code\n",
    "  return (inp.argmax(dim=1)[mask]==targ[mask]).float().mean() \n",
    "opt = ranger\n",
    "learn = unet_learner(dls, resnet101, metrics=acc_camvid, self_attention=True, act_cls=Mish, opt_func=opt).to_fp16()\n",
    "\n",
    "device=\"cuda\"\n",
    "learn.load(f'ecadBlur3Watershed_101')\n",
    "model = learn.model\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cgitb import reset\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader\n",
    "import os\n",
    "from PIL import Image\n",
    "from torch.utils.data import Dataset\n",
    "import numpy as np\n",
    "import skimage as sm\n",
    "import skimage.io\n",
    "from matplotlib import pyplot as plt\n",
    "import tifffile\n",
    "from fastai.vision.all import *\n",
    "from skimage.feature import blob, blob_dog, blob_log, blob_doh\n",
    "from os.path import exists\n",
    "\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "BATCH_SIZE = 1\n",
    "NUM_WORKERS = 2\n",
    "IMAGE_HEIGHT = 512\n",
    "IMAGE_WIDTH = 512\n",
    "PIN_MEMORY = True\n",
    "\n",
    "# util\n",
    "    \n",
    "def main():\n",
    "    \n",
    "    cwd = os.getcwd()\n",
    "    filenames = os.listdir(cwd + \"/dat/ecadBlur3/test_images\")\n",
    "    if \".DS_Store\" in filenames:\n",
    "        filenames.remove(\".DS_Store\")\n",
    "    if \".ipynb_checkpoints\" in filenames:\n",
    "        filenames.remove(\".ipynb_checkpoints\")\n",
    "    filenames.sort()\n",
    "\n",
    "    label=0\n",
    "    for filename in filenames:\n",
    "        print(filename)\n",
    "        out = learn.predict(f'dat/ecadBlur3/test_images/{filename}')[2][1]\n",
    "#         print(out.shape)\n",
    "        out = np.asarray(out*255, \"uint8\")\n",
    "        tifffile.imwrite(f\"dat_output/ecadBlur3_101/pred_{filename}\", out, \n",
    "            imagej=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}