#9.13-实战Kaggle比赛：狗的品种识别(ImageNet Dogs)
import collections
import d2lzh as d2l
import math
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import data as gdata, loss as gloss, model_zoo, nn
import os
import shutil
import time
import zipfile

#9.13.1-获取和整理数据集
#下载数据集
demo = True
data_dir = '../data/kaggle_dog'
if demo:
    zipfiles = ['train_valid_test_tiny.zip']
else:
    zipfiles = ['train.zip', 'test.zip', 'labels.csv.zip']
for f in zipfiles:
    with zipfile.ZipFile(data_dir + '/' + f, 'r') as z:
        z.extractall(data_dir)
#整理数据集
def reorg_train_valid(data_dir, train_dir, input_dir, valid_ratio, idx_label):
    #训练集中数量最少一类的狗的样本数
    min_n_train_per_label = (collections.Counter(idx_label.values()).most_common()[:-2:-1][0][1])
    #验证集中每类狗的样本数
    n_valid_per_label = math.floor(min_n_train_per_label * valid_ratio)
    label_count = {}
    for train_file in os.listdir(os.path.join(data_dir, train_dir)):
        idx = train_file.split('.')[0]
        label = idx_label[idx]
        d2l.mkdir_if_not_exist([data_dir, input_dir, 'train_valid', label])
        shutil.copy(os.path.join(data_dir, train_dir, train_file),
                    os.path.join(data_dir, input_dir, 'train_valid', label))
        if label not in label_count or label_count[label] < n_valid_per_label:
            d2l.mkdir_if_not_exist([data_dir, input_dir, 'valid', label])
            shutil.copy(os.path.join(data_dir, train_dir, train_file),
                        os.path.join(data_dir, input_dir, 'valid', label))
            label_count[label] = label_count.get(label, 0) + 1
        else:
            d2l.mkdir_if_not_exist([data_dir, input_dir, 'train', label])
            shutil.copy(os.path.join(data_dir, train_dir, train_file),
                        os.path.join(data_dir, input_dir, 'train', label))

#读取训练数据标签。切分验证集并整理测试集
def reorg_dog_data(data_dir, label_file, train_dir, test_dir, input_dir, valid_ratio):
    #读取训练数据标签
    with open(os.path.join(data_fir, label_file), 'r') as f:
        #跳过文件头行(栏名称)
        lines = f.readlines()[1:]
        tokens = [l.rstrip().split('.') for l in lines]
        idx_label = dict(((idx, label) for idx, label in tokens))
    reorg_train_valid(data_dir, train_dir, input_dir, valid_ratio, idx_label)
    #整理测试集
    d2l.mkdir_if_not_exist([data_dir, input_dir, 'test', 'unknown'])
    for test_file in os.listdir(os.path.join(data_dir, test_dir)):
        shutil.copy(os.path.join(data_dir, test_dir, test_file),
                    os.path.join(data_dir, input_dir, 'test', 'unknown'))

if demo:
    #注意，此处使用小数据集并将批量大小相应设小。
    #使用Kaggle比赛的完整数据集时可设批量大小为较大整数
    input_dir, batch_size = 'train_valid_test_tiny', 1
else:
    label_file, train_dir, test_dir = 'label.csv', 'train', 'test'
    input_dir, batch_size, valid_ratio = 'train_valid_test', 128, 0.1
    reorg_dog_data(data_dir, label_file, train_dir, test_dir, input_dir, valid_ratio)

#9.13.2-图像增广
transform_train = gdata.vision.transforms.Compose([
    #随机对图像裁剪出面积为圆面积0.08-1倍、且高和宽之比在3/4-4/3的图像，
    #再缩放为高和宽均为224像素的图像
    gdata.vision.transforms.RandomResizedCrop(224, scale=(0.08, 1.0), ratio=(3.0/4.0, 4.0/3.0)),
    gdata.vision.transforms.RandomFlipLeftRight(),
    #随机变化亮度、对比度和饱和度
    gdata.vision.transforms.RandomColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
    #随机加噪声
    gdata.vision.transforms.RandomLighting(0.1),
    gdata.vision.transforms.ToTensor(),
    #对图像的每个通道做标准化
    gdata.vision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] )

#测试时只是用确定性的图像预处理操作
transform_test = gdata.vision.transforms.Compose([
    gdata.vision.transforms.Resize(256),
    #将图像中央的高和宽均为224的正方形区域裁剪出来
    gdata.vision.transforms.CenterCrop(224),
    gdata.vision.transforms.ToTensor(),
    gdata.vision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] )

#读取数据集
train_ds = gdata.vision.ImageFolderDataset(os.path.join(data_dir, input_dir, 'train'), flag=1)
valid_ds = gdata.vision.ImageFolderDataset(os.path.join(data_dir, input_dir, 'valid'), flag=1)
train_valid_ds = gdata.vision.ImageFolderDataset(os.path.join(data_dir, input_dir, 'train_valid'), flag=1)
test_ds = gdata.vision.ImageFolderDataset(os.path.join(data_dir, input_dir, 'test'), flag=1)

train_iter = gdata.DataLoader(train_ds.transform_first(transform_train), batch_size, shuffle=True, last_batch='keep')
valid_iter = gdata.DataLoader(valid_ds.transform_first(transform_test), batch_size, shuffle=True, last_batch='keep')
train_valid_iter = gdata.DataLoader(train_valid_ds.transform_first(transform_train), batch_size, shuffle=True, last_batch='keep')
test_iter = gdata.DataLoader(test_ds.transform_first(transform_test), batch_size, shuffle=False, last_batch='keep')

#9.13.4-定义模型
def get_net(ctx):
    finetune_net = model_zoo.vision.resnet34_v2(pretrained=True)
    #定义新的输出网络
    finetune_net.output_new = nn.HybridSequential(prefix='')
    finetune_net.output_new.add(nn.Dense(256, activation='relu'))
    #120是输出的类别个数
    finetune_net.output_new.add(nn.Dense(120))
    #初始化输出网络
    finetune_net.output_new.initialize(init.Xavier(), ctx=ctx)
    #把模型参数分配到内存或者显存上
    finetune_net.collect_params().reset_ctx(ctx)
    return finetune_net

loss = gloss.SoftmaxCrossEntropyLoss()
def evaluate_loss(data_iter, net, ctx):
    l_sum, n = 0.0, 0
    for X, y in data_iter:
        y = y.as_in_context(ctx)
        output_features = net.features(X.as_in_context(ctx))
        outputs = net.output_new(output_features)
        l_sum += loss(outputs, y).sum().asscalar()
        n += y.size
    return l_sum / n

#9.13.5-定义训练函数
def train(net, train_iter, valid_iter, num_epochs, lr, wd, ctx, lr_period, lr_decay):
    #只训练自定义的小规模输出网络
    trainer = gluon.Trainer(net.output_new.collect_params(), 'sgd', {'learning_rate':lr, 'momentum':0.9, 'wd':wd})
    for epoch in range(num_epochs):
         train_l_sum, n, start = 0.0, 0, time.time()
         if epoch>0 and epoch%lr_perido ==0 :
             trainer.set_learning_rate(trainer.learning_rate*lr_decay)
         for X, y in train_iter:
             y = y.as_in_context(ctx)
             output_features = net.features(X.as_in_context(ctx))
             with autograd.record():
                 outputs = net.output_new(output_features)
                 l = loss(outputs, y).sum()
             l.backward()
             trainer.step(batch_size)
             train_l_sum += l.asscalar()
             n += y.size
         time_s = ("time %.2f sec" % (time.time()-start))
         if valid_iter is not None:
             valid_loss = evaluate_loss(valid_iter, net, ctx)
             epoch_s = ("epoch %d, train loss %f, valid loss %f, " %
                        (epoch + 1, train_l_sum/n, valid_loss))
         else:
             epoch_s = ("epoch %d, train loss %f, " %
                        (epoch + 1, train_l_sum / n))
         print(epoch_s + time_s + ', lr ' + str(trainer.learning_rate))

#9.13.6-训练并验证模型
print('try_gpu')
ctx, num_epochs, lr, wd = d2l.try_gpu(), 1, 0.01, 1e-4
lr_period, lr_decay, net = 10, 0.1, get_net(ctx)
net.hybridize()
print("start train...\n")
train(net, train_iter, valid_iter, num_epochs, lr, wd, ctx, lr_period, lr_decay)
print("end train...\n")

#9.13.7-对测试集分类并在Kaggle提交结果
net = get_net(ctx)
net.hybridize()
print("start valid train...\n")
train(net, train_valid_iter, None, num_epochs, lr, wd, ctx, lr_period, lr_decay)
print("end valid train...\n")
preds = []
for data, label in test_iter:
    output_features = net.features(data.as_in_context(ctx))
    output = nd.softmax(net.output_new(output_features))
    preds.extend(output.asnumpy())
ids = sorted(os.listdir(os.path.join(data_dir, input_dir, 'test/unknown')))
with open('submission.csv', 'w') as f:
    f.write('id,' + ','.join(train_valid_ds.synsets) + '\n')
    for i, output in zip(ids, preds):
        f.write(i.split('.')[0] + ',' + ','.join([str(num) for num in output]) + '\n')
