%matplotlib inline
%reload_ext autoreload
%autoreload 2

from fastai import *
from fastai.vision import *
from fastai.docs import *
from fastai.text import *
torch.backends.cudnn.benchmark=True
import json

import fastText as ft

PATH = Path('../data/imagenet')
TMP_PATH = Path('../data/imagenet/tmp')
TRANS_PATH = Path('../data/translate/')
PATH_TRN = PATH/'train'

ft_vecs = ft.load_model(str((TRANS_PATH/'wiki.en.bin')))

ft_vecs.get_word_vector('king')

np.corrcoef(ft_vecs.get_word_vector('jeremy'), ft_vecs.get_word_vector('Jeremy'))

np.corrcoef(ft_vecs.get_word_vector('banana'), ft_vecs.get_word_vector('Jeremy'))

ft_words = ft_vecs.get_words(include_freq=True)
ft_word_dict = {k:v for k,v in zip(*ft_words)}
ft_words = sorted(ft_word_dict.keys(), key=lambda x: ft_word_dict[x])

CLASSES_FN = 'imagenet_class_index.json'
download_url(f'http://files.fast.ai/models/{CLASSES_FN}', TMP_PATH/CLASSES_FN)

WORDS_FN = 'classids.txt'
download_url(f'http://files.fast.ai/data/{WORDS_FN}', PATH/WORDS_FN)

class_dict = json.load((TMP_PATH/CLASSES_FN).open())
classids_1k = dict(class_dict.values())
nclass = len(class_dict); nclass

class_dict['0']

classid_lines = (PATH/WORDS_FN).open().readlines()
classid_lines[:5]

classids = dict(l.strip().split() for l in classid_lines)
len(classids),len(classids_1k)

lc_vec_d = {w.lower(): ft_vecs.get_word_vector(w) for w in ft_words[-1000000:]}

syn_wv = [(k, lc_vec_d[v.lower()]) for k,v in classids.items()
          if v.lower() in lc_vec_d]
syn_wv_1k = [(k, lc_vec_d[v.lower()]) for k,v in classids_1k.items()
          if v.lower() in lc_vec_d]
syn2wv = dict(syn_wv)
len(syn2wv), len(syn_wv_1k)

pickle.dump(syn2wv, (TMP_PATH/'syn2wv.pkl').open('wb'))
pickle.dump(syn_wv_1k, (TMP_PATH/'syn_wv_1k.pkl').open('wb'))

syn2wv = pickle.load((TMP_PATH/'syn2wv.pkl').open('rb'))
syn_wv_1k = pickle.load((TMP_PATH/'syn_wv_1k.pkl').open('rb'))

images = []
img_vecs = []
images_val = []
img_vecs_val = []

for d in (PATH/'train').iterdir():
    if d.name not in syn2wv: continue
    vec = syn2wv[d.name]
    for f in d.iterdir():
        images.append(str(f.relative_to(PATH)))
        img_vecs.append(vec)

n_val=0
for d in (PATH/'valid').iterdir():
    if d.name not in syn2wv: continue
    vec = syn2wv[d.name]
    for f in d.iterdir():
        images_val.append(str(f.relative_to(PATH)))
        img_vecs_val.append(vec)
        n_val += 1

n_val

img_vecs = np.stack(img_vecs)
img_vecs.shape

pickle.dump(images, (TMP_PATH/'images.pkl').open('wb'))
pickle.dump(img_vecs, (TMP_PATH/'img_vecs.pkl').open('wb'))
pickle.dump(images_val, (TMP_PATH/'images_val.pkl').open('wb'))
pickle.dump(img_vecs_val, (TMP_PATH/'img_vecs)val.pkl').open('wb'))

images = pickle.load((TMP_PATH/'images.pkl').open('rb'))
img_vecs = pickle.load((TMP_PATH/'img_vecs.pkl').open('rb'))
images_val = pickle.load((TMP_PATH/'images_val.pkl').open('rb'))
img_vecs_val = pickle.load((TMP_PATH/'img_vecs_val.pkl').open('rb'))

folder_path = (PATH/"").absolute()
images = [folder_path/image for image in images]
images_val = [folder_path/image_val for image_val in images_val]

n = len(images); n

n_val = len(images_val); n_val

n+n_val

train_ds = ImageDataset(images, img_vecs)
valid_ds = ImageDataset(images_val, img_vecs_val)

train_ds.classes = range(300)
valid_ds.classes = range(300)

tfms = [[flip_lr(), crop_pad(size=224)], [flip_lr(), crop_pad(size=224)]]

data = DataBunch.create(train_ds, valid_ds, path=PATH, device=torch.device('cuda'), ds_tfms = tfms, tfms=imagenet_norm, size=224)

??get_transforms

len(data.train_dl)

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

x[0],y[0]

len(data.valid_dl)

x_val,y_val = next(iter(data.valid_dl))
x_val.shape, y_val.shape

x_val[0],y_val[0]

from PIL import Image

Image.open(x)

import bcolz, threading
from tqdm import tqdm
from torch.utils.data import Dataset

??ConvLearner

learn = ConvLearner(data, tvm.resnet50, ps=[0.2,0.2], lin_ftrs=[1024], pretrained=True, callback_fns=BnFreeze)

body = learn.model[0]

layers = list(body.children())
layers += [AdaptiveConcatPool2d(), Flatten()]   
body = nn.Sequential(*layers)

body

nf = num_features(body)*2

nf

class FCDataset(Dataset):
    def __init__(self, x, y):
        self.x = x
        self.y = y
        
    def __getitem__(self, index):
        return (self.x[index], self.y[index])

    def __len__(self):
        return len(self.x)

learn.opt_fn = partial(AdamW, betas=(0.9,0.99))

def cos_loss(inp,targ): return 1 - F.cosine_similarity(inp,targ).mean()
learn.loss_fn = cos_loss

def get_activations(path, model_name, tmp_path, nf, force=False): 
    tmpl = f'_{model_name}.bc'
    names = [os.path.join(path/'tmp', p+tmpl) for p in ('x_act', 'x_act_val')]
    if os.path.exists(names[0]) and not force:
        activations = [bcolz.open(p) for p in names]
    else:
        activations = [create_empty_bcolz(nf,p) for p in names]
    return activations

def create_empty_bcolz(n, name):
    return bcolz.carray(np.zeros((0,n), np.float32), chunklen=1, mode='w', rootdir=name)

def predict_to_bcolz(m, gen, arr, workers=4):
    arr.trim(len(arr))
    lock=threading.Lock()
    m.eval()
    for x,*_ in tqdm(gen):
        y = to_np(m(x.data).detach())
        with lock:
            arr.append(y)
    arr.flush()

def save_fc1(data, model, path, model_name, tmp_path, nf):
    act, val_act = get_activations(path, model_name, tmp_path, nf)
    m=model
    if len(act)!=len(data.train_ds):
        predict_to_bcolz(m, data.train_dl, act)
    if len(val_act)!=len(data.valid_ds):
        predict_to_bcolz(m, data.valid_dl, val_act)
    
    fc_data = FCDataset(act, img_vecs)
    fc_data_val = FCDataset(val_act, img_vecs_val)
    
    fc_data.classes = data.classes
    fc_data_val.classes = data.classes
    
    fc_db = DataBunch.create(fc_data, fc_data_val, path=PATH, device=torch.device('cuda'), bs=128)
    return fc_db

def num_features(m:Model)->int:
    "Return the number of output features for a `model`."
    for l in reversed(flatten_model(m)):
        if hasattr(l, 'num_features'): 
            return l.num_features

fc_db = save_fc1(data, body, learn.path, 'resnet50', TMP_PATH, nf)

for x,y in iter(fc_db.valid_dl):
        print(x.shape, y.shape)

head = learn.model[1][2:]

head

learn_head = Learner(data=fc_db, model=head, opt_fn = partial(AdamW, betas=(0.9,0.99)), loss_fn = cos_loss)

learn_head.lr_find(start_lr=1e-4, end_lr=1e15)

learn_head.recorder.plot()

lr = 1e-3
wd = 1e-7

lr

learn_head.fit_one_cycle(cyc_len=2, max_lr=lr, wd=wd, div_factor=20, pct_start=0.2)

learn_head.save('pre0')

lrs = np.array([lr/1000,lr/100,lr])

learn.fit_one_cycle(3, lrs, wd=wd)

syns, wvs = list(zip(*syn_wv_1k))
wvs = np.array(wvs)

%time pred_wv = learn.predict()

start=300

denorm = md.val_ds.denorm

def show_img(im, figsize=None, ax=None):
    if not ax: fig,ax = plt.subplots(figsize=figsize)
    ax.imshow(im)
    ax.axis('off')
    return ax

def show_imgs(ims, cols, figsize=None):
    fig,axes = plt.subplots(len(ims)//cols, cols, figsize=figsize)
    for i,ax in enumerate(axes.flat): show_img(ims[i], ax=ax)
    plt.tight_layout()

show_imgs(denorm(md.val_ds[start:start+25][0]), 5, (10,10))

import nmslib

def create_index(a):
    index = nmslib.init(space='angulardist')
    index.addDataPointBatch(a)
    index.createIndex()
    return index

def get_knns(index, vecs):
     return zip(*index.knnQueryBatch(vecs, k=10, num_threads=4))

def get_knn(index, vec): return index.knnQuery(vec, k=10)

nn_wvs = create_index(wvs)

idxs,dists = get_knns(nn_wvs, pred_wv)

[[classids[syns[id]] for id in ids[:3]] for ids in idxs[start:start+10]]

all_syns, all_wvs = list(zip(*syn2wv.items()))
all_wvs = np.array(all_wvs)

nn_allwvs = create_index(all_wvs)

idxs,dists = get_knns(nn_allwvs, pred_wv)

[[classids[all_syns[id]] for id in ids[:3]] for ids in idxs[start:start+10]]

nn_predwv = create_index(pred_wv)

en_vecd = pickle.load(open(TRANS_PATH/'wiki.en.pkl','rb'))

vec = en_vecd['boat']

idxs,dists = get_knn(nn_predwv, vec)
show_imgs([open_image(PATH/md.val_ds.fnames[i]) for i in idxs[:3]], 3, figsize=(9,3));

vec = (en_vecd['engine'] + en_vecd['boat'])/2

idxs,dists = get_knn(nn_predwv, vec)
show_imgs([open_image(PATH/md.val_ds.fnames[i]) for i in idxs[:3]], 3, figsize=(9,3));

vec = (en_vecd['sail'] + en_vecd['boat'])/2

idxs,dists = get_knn(nn_predwv, vec)
show_imgs([open_image(PATH/md.val_ds.fnames[i]) for i in idxs[:3]], 3, figsize=(9,3));

fname = 'valid/n01440764/ILSVRC2012_val_00007197.JPEG'

img = open_image(PATH/fname)

show_img(img);

t_img = md.val_ds.transform(img)
pred = learn.predict_array(t_img[None])

idxs,dists = get_knn(nn_predwv, pred)
show_imgs([open_image(PATH/md.val_ds.fnames[i]) for i in idxs[1:4]], 3, figsize=(9,3));