%reload_ext autoreload
%autoreload 2
%matplotlib inline

#export
from nb_104 import *

path = Path('data/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)

mask.shape,mask.data

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

size=640
bs=4

data = (ImageFileList.from_folder(path_img)
        .label_from_func(get_y_fn)
        .split_by_fname_file('../valid.txt')
        .datasets(SegmentationDataset, classes=codes)
        .transform(get_transforms(), size=size, tfm_y=True)
        .databunch(bs=bs)
        .normalize(imagenet_stats))

x,y = data.train_dl.one_batch()
show_xy_images(data.denorm(x.cpu()),y,rows=2, figsize=(7,7))

x.shape,y.shape

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

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

metrics=acc_camvid
# metrics=accuracy

from fastai.vision.learner import cnn_config

#export
@classmethod
def Learner_create_unet(cls, data:DataBunch, arch:Callable, pretrained:bool=True,
             split_on:Optional[SplitFuncOrIdxList]=None, **kwargs:Any)->None:
    "Build Unet learners."
    meta = cnn_config(arch)
    body = create_body(arch(pretrained), meta['cut'])
    model = models.unet.DynamicUnet(body, n_classes=data.c).cuda()
    learn = cls(data, model, **kwargs)
    learn.split(ifnone(split_on,meta['split']))
    if pretrained: learn.freeze()
    apply_init(model[2], nn.init.kaiming_normal_)
    return learn

Learner.create_unet = Learner_create_unet

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

lr_find(learn)
learn.recorder.plot()

lr = 1e-2

learn.fit_one_cycle(5, slice(lr))

learn.save('stage-1')

learn.load('stage-1')

learn.unfreeze()

lr_find(learn)
learn.recorder.plot()

learn.fit_one_cycle(12, slice(1e-5,lr/5))

learn.save('stage-2')

learn.load('stage-2')

learn.fit_one_cycle(24, slice(1e-5,lr/15))

learn.save('stage-3')

learn.fit_one_cycle(12, slice(1e-6,lr/50), wd=1e-6)