{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Camvid segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai import *\n",
    "from fastai.vision import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = Path('data/camvid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path.ls()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_lbl = path/'labels'\n",
    "path_img = path/'images'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = get_image_files(path_img)\n",
    "fnames[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_lbl.ls()[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f = fnames[0]\n",
    "img = open_image(img_f)\n",
    "img.show(figsize=(5,5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "codes = np.loadtxt(path/'codes.txt', dtype=str)\n",
    "codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_y_fn(fn): return path_lbl/f'{fn.name[:-4]}_P.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = open_mask(get_y_fn(img_f))\n",
    "mask.show(figsize=(5,5), alpha=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_fns = np.loadtxt(path/'valid.txt', dtype=str)\n",
    "valid_fns[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_fns = [path_img/o for o in valid_fns]\n",
    "train_fns = list(set(fnames)-set(valid_fns))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_train_fns = [get_y_fn(o) for o in train_fns]\n",
    "y_valid_fns = [get_y_fn(o) for o in valid_fns]\n",
    "len(train_fns),len(valid_fns),len(y_train_fns),len(y_valid_fns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size=128\n",
    "bs=32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = SegmentationDataset(train_fns, y_train_fns, classes=codes)\n",
    "valid_ds = SegmentationDataset(valid_fns, y_valid_fns, classes=codes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tfms,valid_tfms = get_transforms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tds = DatasetTfm(train_ds, train_tfms, size=size, tfm_y=True)\n",
    "valid_tds = DatasetTfm(valid_ds, valid_tfms, size=size, tfm_y=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = DataBunch.create(train_tds, valid_tds, bs=bs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refactor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- valid set\n",
    "  - path\n",
    "  - fnames\n",
    "  - idxs\n",
    "  - min idx\n",
    "  - rand pct\n",
    "- type of data\n",
    "- source of labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- get filenames\n",
    "- get labels\n",
    "- split data\n",
    "- make datasets\n",
    "- get tfms\n",
    "- (make dls)\n",
    "- (use device)\n",
    "- databunch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ItemList():\n",
    "    \"A collection of items with `__len__` and `__getitem__` with `ndarray` indexing semantics\"\n",
    "    def __init__(self, items:Iterator): self.items = np.array(list(items))\n",
    "    def __len__(self)->int: return len(self.items)\n",
    "    def __getitem__(self,i:int)->Any: self.items[i]\n",
    "    def __repr__(self)->str: return f'{self.__class__.__name__} ({len(self)} items)\\n{self.items}'\n",
    "\n",
    "class PathItemList(ItemList):\n",
    "    def __init__(self, items:Iterator, path:PathOrStr='.'):\n",
    "        super().__init__(items)\n",
    "        self.path = Path(path)\n",
    "    def __repr__(self)->str: return f'{super().__repr__()}\\nPath: {self.path}'\n",
    "\n",
    "def join_path(fname:PathOrStr, path:PathOrStr='.')->Path:\n",
    "    \"`Path(path)/Path(fname)`, `path` defaults to current dir\"\n",
    "    return Path(path)/Path(fname)\n",
    "\n",
    "def join_paths(fnames:FilePathList, path:PathOrStr='.')->Collection[Path]:\n",
    "    path = Path(path)\n",
    "    return [join_path(o,path) for o in fnames]\n",
    "\n",
    "def loadtxt_str(path:PathOrStr)->np.ndarray:\n",
    "    \"Return `ndarray` of `str` of lines of text from `path`\"\n",
    "    return np.loadtxt(str(path), str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _df_to_fns_labels(df:pd.DataFrame, fn_col:int=0, label_col:int=1,\n",
    "                      label_delim:str=None, suffix:Optional[str]=None):\n",
    "    \"\"\"Get image file names in `fn_col` by adding `suffix` and labels in `label_col` from `df`.\n",
    "    If `label_delim` is specified, splits the values in `label_col` accordingly.\n",
    "    \"\"\"\n",
    "    if label_delim:\n",
    "        df.iloc[:,label_col] = list(csv.reader(df.iloc[:,label_col], delimiter=label_delim))\n",
    "    labels = df.iloc[:,label_col].values\n",
    "    fnames = df.iloc[:,fn_col].str.lstrip()\n",
    "    if suffix: fnames = fnames + suffix\n",
    "    return fnames.values, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ImageFileList(PathItemList):\n",
    "    @classmethod\n",
    "    def from_folder(cls, path:PathOrStr='.', check_ext:bool=True, recurse=True)->'ImageFileList':\n",
    "        return cls(get_image_files(path, check_ext=check_ext, recurse=recurse), path)\n",
    "    \n",
    "    def label_from_func(self, func:Callable)->Collection:\n",
    "        return LabelList([(o,func(o)) for o in self.items], self.path)\n",
    "    \n",
    "    def label_from_re(self, pat:str, full_path:bool=False)->Collection:\n",
    "        pat = re.compile(pat)\n",
    "        def _inner(o):\n",
    "            s = str(o if full_path else o.name)\n",
    "            res = pat.search(s)\n",
    "            assert res,f'Failed to find \"{pat}\" in \"{s}\"'\n",
    "            return res.group(1)\n",
    "        return self.label_from_func(_inner)\n",
    "    \n",
    "    def label_from_df(self, df, fn_col:int=0, label_col:int=1, sep:str=None, folder:PathOrStr='.', \n",
    "                      suffix:str=None)->Collection:\n",
    "        fnames, labels = _df_to_fns_labels(df, fn_col, label_col, sep, suffix)\n",
    "        fnames = join_paths(fnames, self.path/Path(folder))\n",
    "        return LabelList([(fn, np.array(lbl, dtype=np.object)) for fn, lbl in zip(fnames, labels) if fn in self.items], \n",
    "                         self.path)\n",
    "    \n",
    "    def label_from_csv(self, csv_fname, header:Optional[Union[int,str]]='infer', fn_col:int=0, label_col:int=1, \n",
    "                       sep:str=None, folder:PathOrStr='.', suffix:str=None)->Collection:\n",
    "        df = pd.read_csv(self.path/csv_fname, header=header)\n",
    "        return self.label_from_df(df, fn_col, label_col, sep, folder, suffix)\n",
    "    \n",
    "    def label_from_folder(self, classes:Collection[str]=None)->Collection:\n",
    "        labels = [fn.parent.parts[-1] for fn in self.items]\n",
    "        if classes is None: classes = uniqueify(labels)\n",
    "        return LabelList([(o,lbl) for o, lbl in zip(self.items, labels) if lbl in classes], self.path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class LabelList(PathItemList):\n",
    "    @property\n",
    "    def files(self): return self.items[:,0]\n",
    "    \n",
    "    def split_by_files(self, valid_fnames:FilePathList)->'SplitData':\n",
    "        valid = [o for o in self.items if o[0] in valid_fnames]\n",
    "        train = [o for o in self.items if o[0] not in valid_fnames]\n",
    "        return SplitData(self.path, LabelList(train), LabelList(valid))\n",
    "    \n",
    "    def split_by_fname_file(self, fname:PathOrStr, path:PathOrStr=None)->'SplitData':\n",
    "        path = Path(ifnone(path, self.path))\n",
    "        fnames = join_paths(loadtxt_str(self.path/fname), path)\n",
    "        return self.split_by_files(fnames)\n",
    "    \n",
    "    def split_by_idx(self, valid_idx:Collection[int])->'SplitData':\n",
    "        valid = [o for i,o in enumerate(self.items) if i in valid_idx]\n",
    "        train = [o for i,o in enumerate(self.items) if i not in valid_idx]\n",
    "        return SplitData(self.path, LabelList(train), LabelList(valid))\n",
    "    \n",
    "    def split_by_folder(self, train:str='train', valid:str='valid')->'SplitData':\n",
    "        n = len(self.path.parts)\n",
    "        folder_name = [o[0].parent.parts[n] for o in self.items]\n",
    "        valid = [o for o in self.items if o[0].parent.parts[n] == valid]\n",
    "        train = [o for o in self.items if o[0].parent.parts[n] == train]\n",
    "        return SplitData(self.path, LabelList(train), LabelList(valid))\n",
    "    \n",
    "    def random_split_by_pct(self, valid_pct:float=0.2)->'SplitData':\n",
    "        rand_idx = np.random.permutation(range(len(self.items)))\n",
    "        cut = int(valid_pct * len(self.items))\n",
    "        return self.split_by_idx(rand_idx[:cut])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@dataclass\n",
    "class SplitData():\n",
    "    path:PathOrStr\n",
    "    train:LabelList\n",
    "    valid:LabelList\n",
    "    test: LabelList=None\n",
    "        \n",
    "    def __post_init__(self): self.path = Path(self.path)\n",
    "    \n",
    "    @property\n",
    "    def lists(self):\n",
    "        res = [self.train,self.valid]\n",
    "        if self.test is not None: res.append(self.test)\n",
    "        return res\n",
    "    \n",
    "    def datasets(self, dataset_cls:type, **kwargs)->'SplitDatasets':\n",
    "        \"Create datasets from the underlying data using `dataset_cls` and passing along the `kwargs`.\"\n",
    "        train = dataset_cls(*self.train.items.T, **kwargs)\n",
    "        dss = [train]\n",
    "        dss += [train.new(*o.items.T, **kwargs) for o in self.lists[1:]]\n",
    "        cls = getattr(train, '__splits_class__', SplitDatasets)\n",
    "        return cls(self.path, *dss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@dataclass\n",
    "class SplitDatasets():\n",
    "    path:PathOrStr\n",
    "    train_ds:Dataset\n",
    "    valid_ds:Dataset\n",
    "    test_ds:Optional[Dataset] = None\n",
    "    \n",
    "    @property\n",
    "    def datasets(self): return [self.train_ds,self.valid_ds]\n",
    "\n",
    "    def transform(self, tfms:TfmList, **kwargs)->'SplitDatasets':\n",
    "        assert not isinstance(self.train_ds, DatasetTfm)\n",
    "        self.train_ds = DatasetTfm(self.train_ds, tfms[0],  **kwargs)\n",
    "        self.valid_ds = DatasetTfm(self.valid_ds, tfms[1],  **kwargs)\n",
    "        if self.test_ds is not None:\n",
    "            self.test_ds = DatasetTfm(self.test_ds, tfms[1],  **kwargs)\n",
    "        return self\n",
    "\n",
    "    def dataloaders(self, **kwargs):\n",
    "        return [DataLoader(o, **kwargs) for o in self.datasets]\n",
    "    \n",
    "    def databunch(self, path=None, **kwargs):\n",
    "        path = Path(ifnone(path, self.path))\n",
    "        return ImageDataBunch.create(*self.datasets, path=path, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = get_transforms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ifl = ImageFileList.from_folder(path_img); ifl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ll = ifl.label_from_func(get_y_fn); ll"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd = ll.split_by_fname_file('../valid.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = get_transforms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dss = sd.datasets(SegmentationDataset, classes=codes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dss.train_ds.classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tdss = dss.transform(tfms, size=128, tfm_y=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = tdss.databunch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = (ImageFileList.from_folder(path_img)\n",
    "        .label_from_func(get_y_fn)\n",
    "        .split_by_fname_file('../valid.txt')\n",
    "        .datasets(SegmentationDataset, classes=codes)\n",
    "        .transform(tfms, size=128, tfm_y=True)\n",
    "        .databunch())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x,y = data.train_dl.one_batch()\n",
    "show_xy_images(x,y,rows=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}