%reload_ext autoreload
%autoreload 2

#export
from nb_005a import *
import pandas as pd
from pandas import Series
import csv
from collections import OrderedDict

PATH = Path('data/planet')
TRAIN = 'train-jpg'
label_csv = PATH/'train_v2.csv'
bs = 64

norm,denorm = normalize_funcs(*imagenet_stats)

df = pd.read_csv(label_csv, header=0)
df.head()

df.iloc[:,1] = list(csv.reader(df.iloc[:,1], delimiter=' '))
df.head()

tst = df.iloc[:,1].apply(lambda x: len(x))
tst.head()

(tst != 1).any()

#export
def uniqueify(x:Series) -> List[Any]: return list(OrderedDict.fromkeys(x).keys())

classes = uniqueify(np.concatenate(df.tags))
' '.join(classes)

class FilesMultiDataset(LabelDataset):
    def __init__(self, fns, labels, classes=None):
        self.classes = ifnone(classes, uniqueify(np.concatenate(labels)))
        self.class2idx = {v:k for k,v in enumerate(self.classes)}
        self.x = np.array(fns)
        self.y = [np.array([self.class2idx[o] for o in l], dtype=np.int64)
                  for l in labels]

    def encode(self, x):
        res = np.zeros((self.c,), np.float32)
        res[x] = 1.
        return res
    
    def get_labels(self, idx): return [self.classes[i] for i in self.y[idx]]
    def __getitem__(self,i): return open_image(self.x[i]), self.encode(self.y[i])
    
    @classmethod
    def from_random_split(cls, path, folder, fns, labels, valid_pct, classes=None, test_name=None):
        train,valid = random_split(valid_pct, f'{path}/{folder}/' + fns, labels)
        train_ds = cls(*train, classes=classes)
        res = [train_ds,cls(*valid, classes=train_ds.classes)]
        if test_name:
            test_fns = get_image_files(Path(path)/test_name)
            test_labels = [[train_ds.classes[0]]]*len(test_fns)
            res.append(cls(test_fns, test_labels, classes=train_ds.classes))
        return res

def get_path(x): return x+'.jpg'

tfms = get_transforms(flip_vert=True, max_lighting=0.1, max_zoom=1.05, max_warp=0.)

np.random.seed(42)
datasets = FilesMultiDataset.from_random_split(
    PATH, TRAIN, get_path(df.image_name), df.tags, valid_pct=0.2, test_name='test-jpg')
valid_ds = datasets[1]

idx=0
valid_ds[idx][0].brightness(0.7).contrast(1.5).show(
    title=valid_ds.get_labels(idx), figsize=(5,5))

valid_ds[1][1]

def get_data(size):
    return DataBunch.create(*datasets, path=PATH, ds_tfms=tfms, size=size,
                            bs=bs, num_workers=12, tfms=norm)

data=get_data(64)

x,y = next(iter(data.valid_dl))
x,y = x.data,y.data
x.shape

#export
def fbeta(y_pred:Tensor, y_true:Tensor, thresh:float=0.5, beta:float=2, eps:float=1e-9, sigmoid:bool=True) -> Rank0Tensor:
    "Computes the f_beta between preds and targets"
    beta2 = beta**2
    if sigmoid: y_pred = y_pred.sigmoid()
    y_pred = (y_pred>thresh).float()
    y_true = y_true.float()
    TP = (y_pred*y_true).sum(dim=1)
    prec = TP/(y_pred.sum(dim=1)+eps)
    rec = TP/(y_true.sum(dim=1)+eps)
    res = (prec*rec)/(prec*beta2+rec+eps)*(1+beta2)
    return res.mean()

def accuracy_thresh(y_pred:Tensor, y_true:Tensor, thresh:float=0.5, sigmoid:bool=True) -> Rank0Tensor:
    if sigmoid: y_pred = y_pred.sigmoid()
    return ((y_pred>thresh)==y_true.byte()).float().mean()

def f2_opt(y_pred, y_true, start=0.15, end=0.25, step=0.01):
    return max([fbeta(y_pred, y_true, th)
                for th in np.arange(start,end,step)])

metrics = [accuracy_thresh, f2_opt]

learn = ConvLearner(data, tvm.resnet34, metrics=metrics, ps=0.4, loss_fn=nn.BCEWithLogitsLoss())

learn.lr_find()
learn.recorder.plot()

def train_size(size, it_num, lr):
    learn.data=get_data(size)
    learn.freeze()
    learn.fit_one_cycle(5, slice(lr))
    learn.save(str(it_num))
    learn.unfreeze()
    learn.fit_one_cycle(5, slice(lr/20, lr/2), pct_start=0.1)
    learn.save(str(it_num+1))

lr=1e-2

train_size(64, 0, lr)

train_size(128, 2, lr/2)

train_size(256, 2, lr/5)

learn.fit_one_cycle(5, slice(lr/100, lr/10), pct_start=0.1)