{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from nb_006 import *\n",
    "import gc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Carvana"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path('data/carvana')\n",
    "PATH_PNG = PATH/'train_masks_png'\n",
    "PATH_X_FULL = PATH/'train'\n",
    "PATH_X_128 = PATH/'train-128'\n",
    "PATH_Y_FULL = PATH_PNG\n",
    "PATH_Y_128 = PATH/'train_masks-128'\n",
    "\n",
    "PATH_X = PATH_X_128\n",
    "PATH_Y = PATH_Y_128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_y_fn(x_fn): return PATH_Y/f'{x_fn.name[:-4]}_mask.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_datasets(path):\n",
    "    x_fns = [o for o in path.iterdir() if o.is_file()]\n",
    "    y_fns = [get_y_fn(o) for o in x_fns]\n",
    "    mask = [o>=1008 for o in range(len(x_fns))]\n",
    "    arrs = arrays_split(mask, x_fns, y_fns)\n",
    "    return [MatchedFilesDataset(*o) for o in arrs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size=128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tfm_datasets(size):\n",
    "    datasets = get_datasets(PATH_X_128 if size<=128 else PATH_X_FULL)\n",
    "    tfms = get_transforms(do_flip=True, max_rotate=4, max_lighting=0.2)\n",
    "    return transform_datasets(*datasets, tfms=tfms, tfm_y=True, size=size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "default_norm,default_denorm = normalize_funcs(*imagenet_stats)\n",
    "bs = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(size, bs):\n",
    "    return DataBunch.create(*get_tfm_datasets(size), bs=bs, tfms=default_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_data(size, bs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hook_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "Sizes = List[List[int]]\n",
    "\n",
    "def in_channels(m:Model) -> List[int]:\n",
    "    \"Returns the shape of the first weight layer\"\n",
    "    for l in flatten_model(m):\n",
    "        if hasattr(l, 'weight'): return l.weight.shape[1]\n",
    "    raise Exception('No weight layer')\n",
    "\n",
    "def model_sizes(m:Model, size:tuple=(256,256), full:bool=True) -> Tuple[Sizes,Tensor,Hooks]:\n",
    "    \"Passes a dummy input through the model to get the various sizes\"\n",
    "    hooks = hook_outputs(m)\n",
    "    ch_in = in_channels(m)\n",
    "    x = torch.zeros(1,ch_in,*size)\n",
    "    x = m.eval()(x)\n",
    "    res = [o.stored.shape for o in hooks]\n",
    "    if not full: hooks.remove()\n",
    "    return res,x,hooks if full else res\n",
    "\n",
    "def get_sfs_idxs(sizes:Sizes, last:bool=True) -> List[int]:\n",
    "    \"Get the indexes of the layers where the size of the activation changes\"\n",
    "    if last:\n",
    "        feature_szs = [size[-1] for size in sizes]\n",
    "        sfs_idxs = list(np.where(np.array(feature_szs[:-1]) != np.array(feature_szs[1:]))[0])\n",
    "        if feature_szs[0] != feature_szs[1]: sfs_idxs = [0] + sfs_idxs\n",
    "    else: sfs_idxs = list(range(len(sfs)))\n",
    "    return sfs_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class UnetBlock(nn.Module):\n",
    "    \"An basic unet block\"\n",
    "    def __init__(self, up_in_c:int, x_in_c:int, hook:Hook):\n",
    "        super().__init__()\n",
    "        self.hook = hook\n",
    "        ni = up_in_c\n",
    "        self.upconv = conv2d_trans(ni, ni//2) # H, W -> 2H, 2W\n",
    "        ni = ni//2 + x_in_c\n",
    "        self.conv1 = conv2d(ni, ni//2)\n",
    "        ni = ni//2\n",
    "        self.conv2 = conv2d(ni, ni)\n",
    "        self.bn = nn.BatchNorm2d(ni)\n",
    "\n",
    "    def forward(self, up_in:Tensor) -> Tensor:\n",
    "        up_out = self.upconv(up_in)\n",
    "        cat_x = torch.cat([up_out, self.hook.stored], dim=1)\n",
    "        x = F.relu(self.conv1(cat_x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        return self.bn(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class DynamicUnet(nn.Sequential):\n",
    "    \"Unet created from a given architecture\"\n",
    "    def __init__(self, encoder:Model, n_classes:int, last:bool=True):\n",
    "        imsize = (256,256)\n",
    "        sfs_szs,x,self.sfs = model_sizes(encoder, size=imsize)\n",
    "        sfs_idxs = reversed(get_sfs_idxs(sfs_szs, last))\n",
    "        \n",
    "        ni = sfs_szs[-1][1]\n",
    "        middle_conv = nn.Sequential(conv2d_relu(ni, ni*2, bn=True), conv2d_relu(ni*2, ni, bn=True))\n",
    "        x = middle_conv(x)\n",
    "        layers = [encoder, nn.ReLU(), middle_conv]\n",
    "\n",
    "        for idx in sfs_idxs:\n",
    "            up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])\n",
    "            unet_block = UnetBlock(up_in_c, x_in_c, self.sfs[idx])\n",
    "            layers.append(unet_block)\n",
    "            x = unet_block(x)\n",
    "\n",
    "        ni = unet_block.conv2.out_channels\n",
    "        if imsize != sfs_szs[0][-2:]: layers.append(conv2d_trans(ni, ni))\n",
    "        layers.append(conv2d(ni, n_classes, 1))\n",
    "        super().__init__(*layers)\n",
    "\n",
    "    def __del__(self):\n",
    "        if hasattr(self, \"sfs\"): self.sfs.remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics=[accuracy_thresh,dice]\n",
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CrossEntropyFlat(nn.CrossEntropyLoss):\n",
    "    def forward(self, input, target):\n",
    "        n,c,*_ = input.shape\n",
    "        return super().forward(input.view(n, c, -1), target.view(n, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "body = create_body(tvm.resnet34(True), 2)\n",
    "model = DynamicUnet(body, n_classes=2).cuda()\n",
    "\n",
    "learn = Learner(data, model, metrics=metrics,\n",
    "                loss_fn=CrossEntropyFlat())\n",
    "\n",
    "learn.split([model[0][6], model[1]])\n",
    "learn.freeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_find(learn)\n",
    "learn.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(1, slice(lr), pct_start=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(6, slice(lr), pct_start=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.save('u0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x,py = learn.pred_batch()\n",
    "\n",
    "for i, ax in enumerate(plt.subplots(4,4,figsize=(10,10))[1].flat):\n",
    "    show_image(default_denorm(x[i].cpu()), py[i]>0, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.unfreeze()\n",
    "lr=1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(6, slice(lr/100,lr), pct_start=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size=512\n",
    "bs = 8\n",
    "learn.data = get_data(size, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.freeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}