{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from nb_006b import *\n",
    "from concurrent.futures import ProcessPoolExecutor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Camvid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path('data/camvid_orig')\n",
    "PATH_X = PATH/'701_StillsRaw_full'\n",
    "PATH_Y = PATH/'LabeledApproved_full'\n",
    "PATH_Y_PROCESSED = PATH/'LabelProcessed'\n",
    "label_csv = PATH/'label_colors.txt'\n",
    "\n",
    "PATH_Y_PROCESSED.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(PATH_Y.iterdir())[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_y_fn(x_fn): return PATH_Y/f'{x_fn.name[:-4]}_L.png'\n",
    "def get_y_proc_fn(y_fn): return PATH_Y_PROCESSED/f'{y_fn.name[:-6]}_P.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_fns = get_image_files(PATH_X)\n",
    "y_fns = [get_y_fn(o) for o in x_fns]\n",
    "y_proc_fns = [get_y_proc_fn(o) for o in y_fns]\n",
    "x_fns[:3],y_fns[:3],y_proc_fns[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_code(l):\n",
    "    a,b = [c for c in l.strip().split(\"\\t\") if c]\n",
    "    return tuple(int(o) for o in a.split(' ')), b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_codes,label_names = zip(*[parse_code(l) for l in open(PATH/\"label_colors.txt\")])\n",
    "label_t = tensor(label_codes)\n",
    "n_labels = len(label_codes)\n",
    "label_codes[:5],label_names[:5], n_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for o in label_names: print(o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name2code = dict(zip(label_names, label_codes))\n",
    "name2id = {v:k for k,v in enumerate(label_names)}\n",
    "void_code = name2id['Void']\n",
    "\n",
    "code2id = ByteTensor(255,255,255).zero_()+void_code\n",
    "for i,code in enumerate(label_codes):\n",
    "    if not code == void_code: code2id[code]=i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def colors_to_codes(color_data):\n",
    "    n = len(color_data)\n",
    "    idxs = tuple(color_data.reshape(n,-1).long())\n",
    "    return code2id[idxs].view(color_data.shape[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "x_img = open_image(x_fns[i])\n",
    "y_img_mask = open_mask(y_fns[i])\n",
    "y_img = Image(y_img_mask.data.int())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_code = colors_to_codes(y_img.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def codes_to_colors(label_data):\n",
    "    h,w = label_data.shape\n",
    "    idxs = label_data.flatten().long()\n",
    "    return Image(label_t.index_select(0, idxs).reshape(h,w,3).permute(2,0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_img2 = codes_to_colors(y_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_img.show(), y_img2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_file(fns):\n",
    "    yfn, pfn = fns\n",
    "    if not pfn.exists():\n",
    "        y_data = open_mask(yfn).px.long()\n",
    "        proc_data = colors_to_codes(y_data)\n",
    "        img = PIL.Image.fromarray(proc_data.numpy())\n",
    "        img.save(pfn)\n",
    "    return pfn\n",
    "\n",
    "def process_label_files(y_fns, y_proc_fns):\n",
    "    ex = ProcessPoolExecutor(16)\n",
    "    for pfn in ex.map(process_file, zip(y_fns, y_proc_fns)):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time process_label_files(y_fns, y_proc_fns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_datasets(path, valid_pct=0.2):\n",
    "    x_fns = get_image_files(path)\n",
    "    y_fns = [get_y_fn(o) for o in x_fns]\n",
    "    y_proc_fns = [get_y_proc_fn(o) for o in y_fns]\n",
    "    total = len(x_fns)\n",
    "    \n",
    "    train, valid = random_split(valid_pct, x_fns, y_proc_fns)\n",
    "    return (MatchedImageDataset(*train), MatchedImageDataset(*valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tfm_datasets(size):\n",
    "    datasets = get_datasets(PATH_X)\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 = 8\n",
    "size = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = get_transforms(do_flip=True, max_rotate=4, max_lighting=0.2)"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = data.train_ds[0]\n",
    "x.shape, y.shape, y.data.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy_no_void(input, target):\n",
    "    target = target.squeeze()\n",
    "    mask = target != void_code\n",
    "    return (input.argmax(dim=1)[mask]==target[mask]).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_no_void(p,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics=[accuracy_no_void]\n",
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "body = create_body(tvm.resnet34(True), 2)\n",
    "model = DynamicUnet(body, n_classes=len(label_codes)).cuda()\n",
    "learn = Learner(data, model, metrics=metrics, loss_fn=CrossEntropyFlat())\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": [
    "lr = 1e-2"
   ]
  },
  {
   "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": [
    "learn.load('u0')\n",
    "x,y = next(iter(learn.data.valid_dl))\n",
    "py = learn.model(x).detach()\n",
    "py = py.softmax(dim=1).max(dim=1, keepdim=True)[1]\n",
    "x,y,py = x.cpu(),y.cpu(),py.cpu()\n",
    "x = default_denorm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 4\n",
    "fig, axs = plt.subplots(n,3,figsize=(10,10), sharey=True)\n",
    "for i in range(n):\n",
    "    Image(x[i]).show(ax=axs[i][0])\n",
    "    codes_to_image(y[i].numpy()).show(ax=axs[i][1])\n",
    "    codes_to_image(py[i].numpy()).show(ax=axs[i][2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.unfreeze()\n",
    "lr=1e-2"
   ]
  },
  {
   "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=640\n",
    "bs = 4\n",
    "learn.data = get_data(size, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learn.freeze()"
   ]
  },
  {
   "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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id2code"
   ]
  },
  {
   "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
}