%reload_ext autoreload
%autoreload 2
%matplotlib inline

from fastai.imports import *
from fastai.torch_imports import *
from fastai.transforms import *
from fastai.model import *
from fastai.dataset import *
from fastai.sgdr import *
from fastai.plots import *
from fastai.conv_learner import *

PATH = "data/dogbreeds/"
sz = 224
arch = resnext101_64
bs = 64

label_csv = f'{PATH}labels.csv'
n = len(list(open(label_csv)))-1
val_idxs = get_cv_idxs(n)

val_idxs, n, len(val_idxs)

"""
Explanation: Lesson 1 Dogbreeds CodeAlong
End of explanation
"""


!ls {PATH}

label_df = pd.read_csv(label_csv)
label_df.head()

# use Pandas to create pivot table which shows how many of each label:
label_df.pivot_table(index='breed', aggfunc=len).sort_values('id', ascending=False)

tfms = tfms_from_model(arch, sz, aug_tfms=transforms_side_on, max_zoom=1.1)
data = ImageClassifierData.from_csv(PATH, folder='train', csv_fname=f'{PATH}labels.csv', 
                                    test_name='test', val_idxs=val_idxs, suffix='.jpg',
                                    tfms=tfms, bs=bs)

fn = PATH + data.trn_ds.fnames[0]; fn

img = PIL.Image.open(fn); img

img.size

size_d = {k: PIL.Image.open(PATH + k).size for k in data.trn_ds.fnames}

row_sz, col_sz = list(zip(*size_d.values()))

row_sz = np.array(row_sz); col_sz = np.array(col_sz)

row_sz[:5]

plt.hist(row_sz);

plt.hist(row_sz[row_sz < 1000])

plt.hist(col_sz);

plt.hist(col_sz[col_sz < 1000])

len(data.trn_ds), len(data.test_ds)

len(data.classes), data.classes[:5]

"""
Explanation: 2. Initial Exploration
End of explanation
"""


def get_data(sz, bs):
    tfms = tfms_from_model(arch, sz, aug_tfms=transforms_side_on, max_zoom=1.1)
    data = ImageClassifierData.from_csv(PATH, 'train', f'{PATH}labels.csv', test_name='test', 
                                        num_workers=4, val_idxs=val_idxs, suffix='.jpg', 
                                        tfms=tfms, bs=bs)
    return data if sz > 300 else data.resize(340, 'tmp')

"""
Explanation: 3. Initial Model
starting w/ small images, large batch sizes to train model v.fast in beginning; increase image size and decrease batch-size as go along.
End of explanation
"""


data = get_data(sz, bs)

learn = ConvLearner.pretrained(arch, data, precompute=True) # GTX870M;bs=64;sz=224;MEM:2431/3017

learn.fit(1e-2, 5)

"""
Explanation: 3.1 Precompute
End of explanation
"""


from sklearn import metrics

# data = get_data(sz, bs)

learn = ConvLearner.pretrained(arch, data, precompute=True, ps=0.5)

learn.fit(1e-2, 2)

lrf = learn.find_lr()
learn.sched.plot()

# turn precompute off then use dataug
learn.precompute = False

learn.fit(1e-2, 5, cycle_len=1)

learn.save('224_pre')

learn.load('224_pre')

"""
Explanation: 3.2 Augment
End of explanation
"""


learn.set_data(get_data(299, bs=32))
learn.freeze() # just making all but last layer already frozen

learn.fit(1e-2, 3, cycle_len=1) # precompute is off so DataAugmentation is back on

learn.fit(1e-2, 3, cycle_len=1, cycle_mult=2)

log_preds, y = learn.TTA()
probs = np.exp(log_preds)
accuracy(log_preds, y), metrics.log_loss(y, probs)

learn.save('299_pre')

# learn.load('299_pre')

learn.fit(1e-2, 1, cycle_len=2)

learn.save('299_pre')

log_preds, y = learn.TTA()
probs = np.exp(log_preds)
accuracy(log_preds, y), metrics.log_loss(y, probs)

SUBM = f'{PATH}subm/'
os.makedirs(SUBM, exist_ok=True)
df.to_csv(f'{SUBM}subm.gz', compression='gzip', index=False)

FileLink(f'{SUBM}subm.gz')

"""
Explanation: 3.3 Increase Size

If you train smth on a smaller size, you can call learn.set_data() and pass in a larger sized dataset. That'll take your model, however it's trained so far, and continue to train on larger images.
This is another way to get SotA results. Starting training on small images for a few epochs, then switching to larger images and continuing training is an amazing effective way to avoid overfitting.

J.Howard (paraphrased)
NOTE: Fully-Convolutional Architectures only.
End of explanation
"""


fn = data.val_ds.fnames[0]

fn

Image.open(PATH+fn).resize((150,150))

trn_tfms, val_tfms = tfms_from_model(arch, sz)

learn = ConvLearner.pretrained(arch, data)
learn.load('299_pre')

# ds = FilesIndexArrayDataset([fn], np.array([0]), val_tfms, PATH)
# dl = DataLoader(ds)
# preds = learn.predict_dl(dl)
# np.argmax(preds)

im = trn_tfms(Image.open(PATH+fn))
preds = to_np(learn.model(V(T(im[None]).cuda())))
np.argmax(preds)

trn_tfms, val_tfms = tfms_from_model(arch, sz)

im = val_tfms(Image.open(PATH+fn)) # or could apply trn_tfms(.)
preds = learn.predict_array(im[None]) # index into image as[None] to create minibatch of 1 img
np.argmax(preds)

"""
Explanation: 6. Individual Prediction
End of explanation
"""


