import pandas as pd
import numpy as np
from sklearn.utils import shuffle
import mxnet as mx
from mxnet import gluon
from mxnet.gluon import nn
from mxnet.image import color_normalize
import mxnet.autograd as autograd
from mxnet.gluon.model_zoo import vision as models

#더미 파일 리스트를 구축한다. , im2rec.py는 https://github.com/apache/incubator-mxnet/blob/master/tools/im2rec.py 
%run im2rec.py --list list catsdogs_full train

#아래 보는것과 같이 텝으로 구분되어서 파일 인덱스, 레이블(클래스), 사진파일명 정보가 구성된다. 
!head catsdogs_full.lst

#pandas 테이블로 파일을 읽어들여 레이블을 마킹하고 랜덤 셔플한 뒤 학습과 테스트셋으로 만든다. 
catdoglist = pd.read_csv('catsdogs_full.lst',sep='\t', names=['idx', 'class', 'fn'])

catdoglist['class'] = [1  if i else 0 for i in catdoglist['fn'].str.contains('dog')]

catdoglist = shuffle(catdoglist)

train = catdoglist.iloc[:6500, ]

train.shape

test = catdoglist.iloc[6500:, ]

train.to_csv('catsdogs_train.lst',doublequote=False,sep='\t', header=False, index=False)
test.to_csv('catsdogs_test.lst',doublequote=False,sep='\t', header=False, index=False)

#실제 Image RecordIO dataset을 만든다. 
!python im2rec.py --num-thread 10 catsdogs_t train

#학습은 GPU에서 
ctx = mx.gpu(0)

class cats_and_dogs(gluon.HybridBlock):
    def __init__(self, num_class, num_hidden , **kwargs):
        super(cats_and_dogs, self).__init__(**kwargs)
        
        with self.name_scope():
            #model zoo에서 vgg16 모형을 가져온다. 
            #이렇게 vgg를 클래스 내부에서 로딩하면 인스턴스마다 context, prefix 가 통일되게 생성되어 모형 저장/로딩시 용이한 측면이 있다.  
            self.vgg_net =  models.vgg16(pretrained=True, root=".models").features
            self.classifier = nn.HybridSequential()
            with self.classifier.name_scope():
                self.classifier.add(nn.Dense(num_hidden, activation="relu"))
                self.classifier.add(nn.Dense(int(num_hidden/2), activation="relu"))
                self.classifier.add(nn.Dense(units=num_class))            
    
    def hybrid_forward(self, F, inputs):
        o_vgg = self.vgg_net(inputs)
        pred = self.classifier(o_vgg)
        return pred
        
        
        

model = cats_and_dogs(num_class=2, num_hidden = 500)

#Symbolic 텐서로 컴파일한다. 
#확실히 GPU메모리 사용량은 줄어드나, 학습 속도가 2배 넘게 빨리지지는 않는거 같다. 
model.hybridize()

print(model)

#생성한 rec 파일로 이터레이터를 생성한다. 
#Tensorflow와 케라스에서와는 다르게 채널이 먼저 나온다.
#이미지 학습을 하는데 ImageRecordIter는 거의 필수라고 생각하는데, 이는 다양한 data augmentation의 옵션을 제공하기 때문이다.
batch_size = 50  # GPU메모리 크기에 따라 정해준다. 
valid_batch_size = 30 #validation은 적어도 큰 문제가 없다. 
train_iter = mx.io.ImageRecordIter(path_imgrec='catsdogs_train.rec',
                                   min_img_size=256,
                                   data_shape=(3, 224, 224),
                                   rand_crop=True,
                                   shuffle=True,
                                   batch_size=batch_size,
                                   max_random_scale=1.5,
                                   min_random_scale=0.75,
                                   rand_mirror=True)
val_iter = mx.io.ImageRecordIter(path_imgrec="catsdogs_test.rec",
                                 min_img_size=256,
                                 data_shape=(3, 224, 224),
                                 batch_size=valid_batch_size)

def evaluate(net, data_iter, ctx):
    data_iter.reset()
    acc = mx.metric.Accuracy()
    for batch in data_iter:
        mx.nd.waitall()
        #reference 
        #https://mxnet.incubator.apache.org/api/python/gluon/model_zoo.html
        data = color_normalize(batch.data[0]/255,
                               mean=mx.nd.array([0.485, 0.456, 0.406]).reshape((1,3,1,1)),
                               std=mx.nd.array([0.229, 0.224, 0.225]).reshape((1,3,1,1)))
        
        data = data.as_in_context(ctx)
        label = batch.label[0].as_in_context(ctx)
        with autograd.predict_mode():
            output = net(data)
        predictions = mx.nd.argmax(output, axis=1)
        acc.update(preds=predictions, labels=label)
        mx.ndarray.waitall()
    return acc.get()[1]

import mxnet.autograd as autograd
from mxnet import gluon
import logging
import os
import time
import numpy as np
logging.basicConfig(level=logging.INFO)

def train(net, train_iter, val_iter, epochs, ctx, need_init=False):
    if need_init:
        #vgg16 모형은 학습을 fix 한다. 
        # https://github.com/apache/incubator-mxnet/issues/1340
        net.vgg_net.collect_params().setattr('grad_req', 'null')
        net.classifier.collect_params().initialize(mx.init.Xavier(), ctx=ctx)
    #변수들이 모두 같은 ctx에 있는지 다시 확인 
    net.collect_params().reset_ctx(ctx)

    trainer = gluon.trainer.Trainer(net.collect_params(), 'rmsprop', optimizer_params={'learning_rate':0.001})
    loss = gluon.loss.SoftmaxCrossEntropyLoss(sparse_label=False)


    val_accs = evaluate(net, val_iter, ctx)
    logging.info('[Initial] validation accuracy : {}'.format(val_accs))
    tr_loss = []
    val_accs = []
    for epoch in range(epochs):
        tic = time.time()
        train_iter.reset()
        btic = time.time()
        loss_seq = []
        for i, batch in enumerate(train_iter):
            #속도를 좀 감소시키더라도 GPU 메모리에 대한 관리를 더 한다.  
            mx.nd.waitall()
            data = color_normalize(batch.data[0]/255,
                                   mean=mx.nd.array([0.485, 0.456, 0.406]).reshape((1,3,1,1)),
                                   std=mx.nd.array([0.229, 0.224, 0.225]).reshape((1,3,1,1)))
            data = data.as_in_context(ctx)
            label = mx.nd.one_hot(batch.label[0].as_in_context(ctx),2)
            with autograd.record():
                z = net(data)
                L = loss(z, label)
                L.backward()
            trainer.step(data.shape[0])
            btic = time.time()
            curr_loss = mx.nd.mean(L).asscalar()
            loss_seq.append(curr_loss)
            mx.ndarray.waitall()
        logging.info('[Epoch {}] training loss : {}'.format(epoch, np.mean(loss_seq)))
        logging.info('[Epoch %d] time cost: %f'%(epoch, time.time()-tic))
        val_acc = evaluate(net, val_iter, ctx)
        logging.info('[Epoch %d] validation accuracy : %s'%(epoch, val_acc))
        val_accs.append(val_acc)
        tr_loss.append(np.mean(loss_seq))
    return(val_accs, tr_loss)

accs, losses = train(model,train_iter, val_iter, 10, ctx=ctx, need_init=True)

from skimage.color import rgba2rgb
from skimage import io
from matplotlib import pyplot as plt
%matplotlib inline

#cpu에서 동작 시킨다.  
model.collect_params().reset_ctx(mx.cpu(0))

#reference : http://gluon.mxnet.io/chapter08_computer-vision/fine-tuning.html
def classify_dogcat(net, url):
    I = io.imread(url)
    if I.shape[2] == 4:
        I = rgba2rgb(I)
    image = mx.nd.array(I).astype(np.uint8)
    plt.subplot(1, 2, 1)
    plt.imshow(image.asnumpy())
    image = mx.image.resize_short(image, 256)
    image, _ = mx.image.center_crop(image, (224, 224))
    plt.subplot(1, 2, 2)
    plt.imshow(image.asnumpy())
    image = mx.image.color_normalize(image.astype(np.float32)/255,
                                     mean=mx.nd.array([0.485, 0.456, 0.406]),
                                     std=mx.nd.array([0.229, 0.224, 0.225]))
    image = mx.nd.transpose(image.astype('float32'), (2,1,0))
    image = mx.nd.expand_dims(image, axis=0)
    out = mx.nd.SoftmaxActivation(net(image))
    print('Probabilities are: '+str(out[0].asnumpy()))
    result = np.argmax(out.asnumpy())
    outstring = ['cat!', 'dog!']
    print(outstring[result])

%timeit
classify_dogcat(net=model, url="cats.jpg")

%timeit
classify_dogcat(net=model, url="Cute-dog-listening-to-music-1_1.jpg")

model.save_params("cats_and_dogs_model.mdl")

new_model = cats_and_dogs(num_class=2, num_hidden = 500)

new_model.load_params('cats_and_dogs_model.mdl', ctx=mx.cpu(0))

%timeit
classify_dogcat(net=new_model, url="Cute-dog-listening-to-music-1_1.jpg")