%reload_ext autoreload
%autoreload 2
%matplotlib inline

from fastai import *
from fastai.vision import *

path = untar_data(URLs.CAMVID)
path.ls()

path_lbl = path/'labels'
path_img = path/'images'

fnames = get_image_files(path_img)
fnames[:3]

lbl_names = get_image_files(path_lbl)
lbl_names[:3]

img_f = fnames[0]
img = open_image(img_f)
img.show(figsize=(5,5))

get_y_fn = lambda x: path_lbl/f'{x.stem}_P{x.suffix}'

mask = open_mask(get_y_fn(img_f))
mask.show(figsize=(5,5), alpha=1)

src_size = np.array(mask.shape[1:])
src_size,mask.data

codes = np.loadtxt(path/'codes.txt', dtype=str); codes

size = src_size//2
bs=8

src = (SegmentationItemList.from_folder(path_img)
       .split_by_fname_file('../valid.txt')
       .label_from_func(get_y_fn, classes=codes))

data = (src.transform(get_transforms(), size=size, tfm_y=True)
        .databunch(bs=bs)
        .normalize(imagenet_stats))

data.show_batch(2, figsize=(10,7))

data.show_batch(2, figsize=(10,7), ds_type=DatasetType.Valid)

name2id = {v:k for k,v in enumerate(codes)}
void_code = name2id['Void']

def acc_camvid(input, target):
    target = target.squeeze(1)
    mask = target != void_code
    return (input.argmax(dim=1)[mask]==target[mask]).float().mean()

metrics=acc_camvid
# metrics=accuracy

learn = Learner.create_unet(data, models.resnet34, metrics=metrics)

lr_find(learn)
learn.recorder.plot()

lr=1e-2

learn.fit_one_cycle(10, slice(lr))

learn.save('stage-1')

learn.load('stage-1');

learn.show_results()

learn.unfreeze()

lr_find(learn)
learn.recorder.plot()

lrs = slice(1e-5,lr/5)

learn.fit_one_cycle(12, lrs)

learn.recorder.plot_losses()

learn.recorder.plot_lr()

learn.save('stage-2');

size = src_size
bs=4

data = (src.transform(get_transforms(), size=size, tfm_y=True)
        .databunch(bs=bs)
        .normalize(imagenet_stats))

learn = Learner.create_unet(data, models.resnet34, metrics=metrics)

learn.load('stage-2');

lr_find(learn)
learn.recorder.plot()

lr=1e-3

learn.fit_one_cycle(10, slice(lr))

learn.save('stage-1-big')

learn.load('stage-1-big');

learn.unfreeze()

lrs = slice(1e-6,lr)

learn.fit_one_cycle(10, lrs, wd=1e-3)

learn.save('stage-2-big')

learn.load('stage-2-big')

learn.show_results()