#Run once per session
!pip install fastai -q --upgrade

from fastai.vision.all import *

planet_source = untar_data(URLs.PLANET_SAMPLE)
df = pd.read_csv(planet_source/'labels.csv')

df.head()

batch_tfms = aug_transforms(flip_vert=True, max_lighting=0.1, max_zoom=1.05, max_warp=0.)

planet = DataBlock(blocks=(ImageBlock, MultiCategoryBlock),
                   get_x=ColReader(0, pref=f'{planet_source}/train/', suff='.jpg'),
                   splitter=RandomSplitter(),
                   get_y=ColReader(1, label_delim=' '),
                   batch_tfms = batch_tfms)

dls = planet.dataloaders(df)

dls.show_batch(max_n=9, figsize=(12,9))

blocks = (ImageBlock, MultiCategoryBlock)

get_x = lambda x:planet_source/'train'/f'{x[0]}.jpg'

val = df.values[0]; val

get_x(df.values[0])

get_y = lambda x:x[1].split(' ')

planet = DataBlock(blocks=blocks,
                   get_x=get_x,
                   splitter=RandomSplitter(),
                   get_y=get_y,
                   batch_tfms=batch_tfms)

dls = planet.dataloaders(df)
dls.show_batch(max_n=9, figsize=(12,9))

def _planet_items(x): return (
    f'{planet_source}/train/'+x.image_name+'.jpg', x.tags.str.split())

planet = DataBlock.from_columns(blocks=(ImageBlock, MultiCategoryBlock),
                   get_items = _planet_items, 
                   splitter=RandomSplitter(),
                   batch_tfms=batch_tfms)

dls = planet.dataloaders(df)
dls.show_batch(max_n=9, figsize=(12,9))

from torchvision.models import resnet34

from fastai.metrics import accuracy_multi

learn = cnn_learner(dls, resnet34, pretrained=True, metrics=[accuracy_multi])

class BCEWithLogitsLossFlat(BaseLoss):
    "Same as `nn.CrossEntropyLoss`, but flattens input and target."
    def __init__(self, *args, axis=-1, floatify=True, thresh=0.5, **kwargs):
        super().__init__(nn.BCEWithLogitsLoss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)
        self.thresh = thresh

    def decodes(self, x):    return x>self.thresh
    def activation(self, x): return torch.sigmoid(x)

learn.loss_func = BCEWithLogitsLossFlat()

learn.lr_find()

lr = 1e-2
learn = learn.to_fp16()

learn.fit_one_cycle(5, slice(lr))

learn.save('stage-1')

learn.unfreeze()
learn.lr_find()

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

learn.show_results(figsize=(15,15))