# coding=utf8
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""ResNet Train/Eval module.
"""
import time
import sys
import os
# import glob
import numpy as np
import dataloader
import json
from tqdm import tqdm
from collections import Counter
import densenet
import resnet
from PIL import Image

import torchvision

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torch.nn.functional as F

from sklearn.metrics import roc_auc_score

from tools import parse
from glob import glob
from skimage import measure
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
reload(sys)
sys.setdefaultencoding('utf8')
import traceback
# from moxing.framework import file

args = parse.args
# anchor大小
args.anchors = [8, 12, 18, 27, 40, 60]
args.stride = 8
args.image_size = [1300,64]
#args.image_size = [64, 64]
datadir = parse.datadir
#datadir = '/home/work/user-job-dir/ocr_densenet/data'
data_dir_obs = 's3://densenet-214/data/dataset'

model_dir='/home/work/user-job-dir/model_15'
model_dir_obs='s3://densenet-214/out/model_15'
###事先执行#####
#print 'datadir is exist ? ', os.path.exists(datadir)
print 'datadir path', args.data_dir
print 'data_dir_obs path', data_dir_obs
print 'current path', os.getcwd()
print '================================================'
print os.listdir(os.getcwd())
# file.copy_parallel(data_dir_obs, args.data_dir)
# file.copy_parallel(model_dir_obs, model_dir)
print 'batch-size:',args.batch_size
print '===================data_dir============================='
# print os.listdir(args.data_dir)
class DenseNet121(nn.Module):
    """Model modified.

    The architecture of our model is the same as standard DenseNet121
    except the classifier layer which has an additional sigmoid function.

    """

    def __init__(self, out_size):
        super(DenseNet121, self).__init__()
        self.inplanes = 1024
        self.densenet121 = densenet.densenet121(pretrained=False, small=args.small)
        num_ftrs = self.densenet121.classifier.in_features
        self.classifier_font = nn.Sequential(
            # 这里可以用fc做分类
            # nn.Linear(num_ftrs, out_size)
            # 这里可以用1×1卷积做分类
            nn.Conv2d(num_ftrs, out_size, kernel_size=1, bias=False)
        )
        self.train_params = []
        self.unpool = nn.MaxUnpool2d(kernel_size=2, stride=2)

    # 用于构建Resnet中的4个blocks
    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    # 定义数据在层之间的流动顺序
    def forward(self, x, phase='train'):
        feats = self.densenet121(x)  # (32, 1024, 2, 16)
        if not args.small:
            feats = F.max_pool2d(feats, kernel_size=2, stride=2)  # (32, 1024, 1, 8)
        out = self.classifier_font(feats)  # (32, 1824, 1, 8)
        out_size = out.size()
        # print out.size()
        out = out.view(out.size(0), out.size(1), -1)  # (32, 1824, 8)
        # print out.size()
        if phase == 'train':
            out = F.adaptive_max_pool1d(out, output_size=(1)).view(out.size(0), -1)  # (32, 1824)
            return out
        else:
            out = out.transpose(1, 2).contiguous()
            out = out.view(out_size[0], out_size[2], out_size[3], out_size[1])  # (32, 1, 8, 1824)
            return out, feats


class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()
        self.classify_loss = nn.BCELoss()
        self.sigmoid = nn.Sigmoid()
        self.regress_loss = nn.SmoothL1Loss()

    def forward(self, font_output, font_target, weight=None, use_hard_mining=False):
        font_output = self.sigmoid(font_output)
        font_loss = F.binary_cross_entropy(font_output, font_target, weight)

        # hard_mining 
        if use_hard_mining:
            font_output = font_output.view(-1)
            font_target = font_target.view(-1)
            pos_index = font_target > 0.5
            neg_index = font_target == 0

            # pos
            pos_output = font_output[pos_index]
            pos_target = font_target[pos_index]
            num_hard_pos = max(len(pos_output) / 4, min(5, len(pos_output)))
            if len(pos_output) > 5:
                pos_output, pos_target = hard_mining(pos_output, pos_target, num_hard_pos, largest=False)
            pos_loss = self.classify_loss(pos_output, pos_target) * 0.5

            # neg
            num_hard_neg = len(pos_output) * 2
            neg_output = font_output[neg_index]
            neg_target = font_target[neg_index]
            neg_output, neg_target = hard_mining(neg_output, neg_target, num_hard_neg, largest=True)
            neg_loss = self.classify_loss(neg_output, neg_target) * 0.5

            font_loss += pos_loss + neg_loss

        else:
            pos_loss, neg_loss = font_loss, font_loss
        return [font_loss, pos_loss, neg_loss]

    def _forward(self, font_output, font_target, weight, bbox_output=None, bbox_label=None, seg_output=None,
                 seg_labels=None):
        font_output = self.sigmoid(font_output)
        font_loss = F.binary_cross_entropy(font_output, font_target, weight)

        acc = []
        if bbox_output is not None:
            # bbox_loss = 0
            bbox_output = bbox_output.view((-1, 4))
            bbox_label = bbox_label.view((-1, 4))
            pos_index = bbox_label[:, -1] >= 0.5
            pos_index = pos_index.unsqueeze(1).expand(pos_index.size(0), 4)
            neg_index = bbox_label[:, -1] <= -0.5
            neg_index = neg_index.unsqueeze(1).expand(neg_index.size(0), 4)

            # 正例
            pos_label = bbox_label[pos_index].view((-1, 4))
            pos_output = bbox_output[pos_index].view((-1, 4))
            lx, ly, ld, lc = pos_label[:, 0], pos_label[:, 1], pos_label[:, 2], pos_label[:, 3]
            ox, oy, od, oc = pos_output[:, 0], pos_output[:, 1], pos_output[:, 2], pos_output[:, 3]
            regress_loss = [
                self.regress_loss(ox, lx),
                self.regress_loss(oy, ly),
                self.regress_loss(od, ld),
            ]
            pc = self.sigmoid(oc)
            acc.append((pc >= 0.5).data.cpu().numpy().astype(np.float32).sum())
            acc.append(len(pc))
            # print pc.size(), lc.size()
            classify_loss = self.classify_loss(pc, lc) * 0.5

            # 负例
            neg_label = bbox_label[neg_index].view((-1, 4))
            neg_output = bbox_output[neg_index].view((-1, 4))
            lc = neg_label[:, 3]
            oc = neg_output[:, 3]
            pc = self.sigmoid(oc)
            acc.append((pc <= 0.5).data.cpu().numpy().astype(np.float32).sum())
            acc.append(len(pc))
            # print pc.size(), lc.size()
            classify_loss += self.classify_loss(pc, lc + 1) * 0.5

            # seg_loss
            seg_output = seg_output.view(-1)
            seg_labels = seg_labels.view(-1)
            pos_index = seg_labels > 0.5
            neg_index = seg_labels < 0.5
            seg_loss = 0.5 * self.classify_loss(seg_output[pos_index], seg_labels[pos_index]) + \
                       0.5 * self.classify_loss(seg_output[neg_index], seg_labels[neg_index])
            seg_tpr = (seg_output[pos_index] > 0.5).data.cpu().numpy().astype(np.float32).sum() / len(
                seg_labels[pos_index])
            seg_tnr = (seg_output[neg_index] < 0.5).data.cpu().numpy().astype(np.float32).sum() / len(
                seg_labels[neg_index])
            # print seg_output[neg_index]
            # print seg_labels[neg_index]




        else:
            return font_loss

        if args.model == 'resnet':
            loss = font_loss + classify_loss + seg_loss
        else:
            loss = font_loss + classify_loss + seg_loss
        for reg in regress_loss:
            loss += reg
        # if args.model == 'resnet':
        #     loss = seg_loss

        return [loss, font_loss, seg_loss, classify_loss] + regress_loss + acc + [seg_tpr, seg_tnr]

        font_num = font_target.sum(0).data.cpu().numpy()
        font_loss = 0
        for di in range(font_num.shape[0]):
            if font_num[di] > 0:
                font_output_i = font_output[:, di]
                font_target_i = font_target[:, di]
                pos_font_index = font_target_i > 0.5
                font_loss += 0.5 * self.classify_loss(font_output_i[pos_font_index], font_target_i[pos_font_index])
                neg_font_index = font_target_i < 0.5
                if len(font_target_i[neg_font_index]) > 0:
                    font_loss += 0.5 * self.classify_loss(font_output_i[neg_font_index], font_target_i[neg_font_index])
        font_loss = font_loss / (font_num > 0).sum()

        return font_loss
        # '''


def hard_mining(neg_output, neg_labels, num_hard, largest=True):
    num_hard = min(max(num_hard, 10), len(neg_output))
    _, idcs = torch.topk(neg_output, min(num_hard, len(neg_output)), largest=largest)
    neg_output = torch.index_select(neg_output, 0, idcs)
    neg_labels = torch.index_select(neg_labels, 0, idcs)
    return neg_output, neg_labels


def save_model(save_dir, phase, name, epoch, f1score, model):
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    save_dir = os.path.join(save_dir, args.model)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    save_dir = os.path.join(save_dir, phase)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    state_dict = model.state_dict()
    for key in state_dict.keys():
        state_dict[key] = state_dict[key].cpu()
    state_dict_all = {
        'state_dict': state_dict,
        'epoch': epoch,
        'f1score': f1score,
    }
    saveStr = '{:s}.ckpt'.format(name)
    torch.save(state_dict_all, os.path.join(save_dir, saveStr))
    # file.copy(os.path.join(save_dir, saveStr), os.path.join(args.save_dir_obs, saveStr))
    if 'best' in name and f1score > 0.3:
        bestStr = '{:s}_{:s}.ckpt'.format(name, str(epoch))
        torch.save(state_dict_all, os.path.join(save_dir, bestStr))
        # file.copy(os.path.join(save_dir, bestStr),
        #           os.path.join(args.save_dir_obs, bestStr))


def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)


def test(model, train_loader, phase='test'):
    print '\ntest {:s}_files, epoch: {:d}'.format(phase, 1)
    mkdir('../../../intermediate_file/recognize_result')
    model.eval()
    f1score_list = []
    recall_list = []
    precision_list = []
    word_index_dict = json.load(open(args.word_index_json, 'r'))
    index_word_dict = {v: k for k, v in word_index_dict.items()}
    # result_file = file.File(datadir + '/data/result/{:d}_{:s}_result.csv'.format(epoch, phase), 'w')
    #modify here for single word
    # result_file = open(datadir + '/data/result/{:d}_{:s}_result_single_64.csv'.format(1, phase), 'w')
    #origin_csv is here
    result_file = open('../../../intermediate_file/recognize_result/result.csv', 'w')
    
    result_file.write('name,content\n')
    name_f1score_dict = dict()

    # 保存densenet生成的feature
    # feat_dir = args.data_dir.replace('dataset', 'feats')
    # mkdir(feat_dir)
    # feat_dir = os.path.join(feat_dir, phase)
    # print feat_dir
    # mkdir(feat_dir)

    names = []
    probs_all = []
    for i, data in enumerate(train_loader):
        if i % 50 == 0:
            print('step[{:d}] OK...'.format(i))
        name = data[0][0].split('/')[-1].split('.seg')[0]
        names.append(name)
        images, labels = [Variable(x.cuda(async=True)) for x in data[1:3]]
        # print 'images.shape:',images.shape,'len(images.size()):',len(images.size())
        if len(images.size()) == 5:
            images = images[0]

        probs, feats = model(images, 'test')
        probs_all.append(probs.data.cpu().numpy().max(2).max(1).max(0))
        # pdb.set_trace()
        probs_temp=probs.data.cpu().numpy()
        #====================================
        probs_temp = probs_temp.max(1) 
        max_preds_line=np.max(probs_temp,axis=2)#max preds among 9115 words
        max_preds_index=np.argmax(probs_temp , axis=2)
        # print 'max_preds_index.shape',max_preds_index.shape
        # print 'max_preds_line.shape',
        # print 'max_preds_index:',max_preds_index.shape,'\nmax_preds_line:',max_preds_line
        probs_temp=probs_temp*0-float('inf')
        # max_preds_index[0,]
        for i in range(max_preds_index.shape[1]):
            probs_temp[0,i,max_preds_index[0,i]]=max_preds_line[0,i]

        preds = probs_temp > 0.0  # (-1, 8, 1824)

        # result_file.write(name+',')
        result = u''
        # last_set = set()
        # all_set = set()
        # if args.feat:
        #     # 保存所有的feat
        #     feats = feats.data.cpu().numpy()
        #     if i == 0:
        #         print feats.shape
        #     np.save(os.path.join(feat_dir, name.replace('.jpg', '.npy')), feats)
        #     if len(feats) > 1:  # feats: [-1, 1024, 1, 8]
        #         # 多个patch
        #         new_feats = []
        #         for i, feat in enumerate(feats):
        #             if i == 0:
        #                 # 第一个patch,保存前6个
        #                 new_feats.append(feat[:, :, :6])
        #             elif i == len(feats) - 1:
        #                 # 最后一个patch,保存后6个
        #                 new_feats.append(feat[:, :, 2:])
        #             else:
        #                 # 保存中间4个
        #                 new_feats.append(feat[:, :, 2:6])
        #         feats = np.concatenate(new_feats, 2)

        # 这种方法用于检测不同区域的同一个字，当同一个字同一个区域出现时，可能检测不到多次
        # print 'name',name,'before preds.shape:',preds.shape,'type(preds):',type(preds) 
             
        # preds = preds.max(1)  # 沿着竖直方向pooling

        
        # print 'max_preds_index:',max_preds_index.shape,'\nmax_preds_line:',max_preds_line

        # print preds
        # if len(preds) > 1:
        #     print name
        # print 'name',name,'after preds.shape:',preds.shape,'type(preds):',type(preds)
    


        lines_words_indexs=[]
        # lines_list=[]
        # print 'name',name
        for patch_i, patch_pred in enumerate(preds):#patch_pred stand for per image feature map 64*9115

            # print 'patch_i:',patch_i
            # each_feature_map_words=set()
            num=0
            for part_i, part_pred in enumerate(patch_pred):#part_i:0-64  part_pred:1*9115
                # print 'name',name,'part_pred.shape:',part_pred.shape,'type(part_pred):',type(part_pred)
                # print part_i
                # new_set = set()
                # print result
                
                for idx, p in enumerate(part_pred):#idx:0-9114
                    if p:
                        w=''
                        # w = index_word_dict[idx]
                        lines_words_indexs.append([part_i,w,idx])
                        # print part_i,':',w
 

                #         new_set.add(w)
                #         # if word_temp!=w
                #         # count+=1
                #         if w not in all_set:
                #             # 从没见过的字
                #             all_set.add(w)
                #             result += w
                #             # count_list.append(count)
                #             # count=0
                #         # elif w!=result[-1]:
                #         #     result += w
                #             # print 'result[-1]:',result[-1]
                #         elif w not in last_set:
                #             # not in last line
                #             if patch_i == 0: #patch always 0
                #                 # 第一个patch # 上一个部分没有这个字
                #                 result += w
                #             elif part_i >= preds.shape[1] / 2:
                #                 # 后续patch的后一半，不写 # 上一个部分没有这个字
                #                 result += w

                # last_set = new_set
        words_total_list=[]
        word_list=[]

        lines_list=[]
        line_list=[]

        index_list=[]
        indexes_list=[]
        neigbour=2
        for i,line_word_index in enumerate(lines_words_indexs):
            if i==0:
                line_list.append(line_word_index[0])
                word_list.append(line_word_index[1])
                index_list.append(line_word_index[2])
            elif (line_word_index[0]-line_list[-1])<neigbour:
                line_list.append(line_word_index[0])
                word_list.append(line_word_index[1])
                index_list.append(line_word_index[2])
            elif (line_word_index[0]-line_list[-1])>=neigbour:
                lines_list.append(line_list)
                words_total_list.append(word_list)
                indexes_list.append(index_list)
                line_list=[]
                word_list=[]
                index_list=[]
                line_list.append(line_word_index[0])
                word_list.append(line_word_index[1])
                index_list.append(line_word_index[2])#clustering according the line number
            if i==(len(lines_words_indexs)-1):
                lines_list.append(line_list)
                words_total_list.append(word_list)
                indexes_list.append(index_list)

        # for i,word in enumerate(words_total_list):#find the most common appear word in each small cluster word list
        #     c=Counter(word)
        #     x=c.most_common()
        #     # print 'x',x
        #     word=[item[0] for item in x if item[1]==x[0][1]]
        #     words_total_list[i]=word#

        for i,index in enumerate(indexes_list):#find the most common appear word_index in each small cluster index list
            c=Counter(index)
            x=c.most_common()
            # print 'x',x
            index=[item[0] for item in x if item[1]==x[0][1]]
            indexes_list[i]=index
        
        for i,indexes in enumerate(indexes_list):#find the biggest probability word index in each small cluster index list
            if len(indexes)>1:
                # print 'lines_list[i]:',lines_list[i]
                # print probs_temp.shape,type(probs_temp)
                prob_temp=np.max(probs_temp[0,lines_list[i],:],0)#become 1*9105
                result_word_index = np.argmax(prob_temp,0)
                indexes=[]
                indexes.append(result_word_index)
                indexes_list[i]=indexes
        for index in indexes_list:#map the word index to the word result
            result+=index_word_dict[index[0]]



            #     probs_temp[lines_list[i]]
            # c=Counter(index)
            # x=c.most_common()
            # # print 'x',x
            # index=[item[0] for item in x if item[1]==x[0][1]]
            # indexes_list[i]=index

            # word=[]
            

        # print 'lines_list:',lines_list



        # print lines_list
        # print json.dumps(words_total_list, encoding="UTF-8", ensure_ascii=False)
        # print indexes_list


        result = result.replace(u'"', u'')
        if u',' in result:
            result = '"' + result + '"'
        if len(result) == 0:
            global_prob = probs.data.cpu().numpy().max(0).max(0).max(0)
            max_index = global_prob.argmax()
            result = index_word_dict[max_index] 
        print 'name',name, 'result',result
        result_file.write(name + ',' + result + '\n')
        # result_file.write('\n')

        if phase == 'test':
            continue

    result_file.close()
    # import pandas as pd
    # re = pd.read_csv(datadir + '/data/result/{:d}_{:s}_result.csv'.format(epoch, phase))
    # re.columns = ['target_file', 'text']
    # submit = pd.read_csv(datadir + '/submission.csv')
    # submit = pd.merge(submit, re, how='left', on=['target_file'])
    # submit = submit.drop(['target_file'], axis=1)
    # submit = submit.replace(to_replace='None', value=20)
    # submit = submit.fillna('上')
    # submit.to_csv(datadir + '/predict.csv', header=True, index=None, encoding='utf-8')
    # file.copy(datadir + '/predict.csv', args.data_dir_obs + '/predict.csv')


def get_weight(labels):
    labels = labels.data.cpu().numpy()
    weights = np.zeros_like(labels)
    # weight_false = 1.0 / ((labels<0.5).sum() + 10e-20)
    # weight_true  = 1.0 / ((labels>0.5).sum() + 10e-20)
    weight_false = 1.0 / ((labels < 0.5).sum(0) + 10e-20)
    label_true = (labels > 0.5).sum(0)
    for i in range(labels.shape[1]):
        label_i = labels[:, i]
        weight_i = np.ones(labels.shape[0]) * weight_false[i]
        # weight_i = np.ones(labels.shape[0]) * weight_false
        if label_true[i] > 0:
            weight_i[label_i > 0.5] = 1.0 / label_true[i]
        weights[:, i] = weight_i
    weights *= np.ones_like(labels).sum() / (weights.sum() + 10e-20)
    weights[labels < -0.5] = 0
    return weights


def train_eval(epoch, model, train_loader, loss, optimizer, best_f1score=0, phase='train'):
    print '\n', epoch, phase
    if 'train' in phase:
        model.train()
    else:
        model.eval()
    loss_list = []
    f1score_list = []
    recall_list = []
    precision_list = []
    for i, data in enumerate(train_loader):
        images, labels = [Variable(x.cuda(async=True)) for x in data[1:3]]
        weights = torch.from_numpy(get_weight(labels)).cuda(async=True)
        probs = model(images)

        # 训练阶段
        if 'train' in phase:
            loss_output = loss(probs, labels, weights, args.hard_mining)
            try:
                optimizer.zero_grad()
                loss_output[0].backward()
                optimizer.step()
                loss_list.append([x.data.cpu().numpy() for x in loss_output])
            except:
                # pass
                traceback.print_exc()
        # 计算 f1score, recall, precision
        '''
        x = probs.data.cpu().numpy() 
        l = labels.data.cpu().numpy()
        print (get_weight(labels) * l).sum()
        l = 1 - l
        print (get_weight(labels) * l).sum()
        print x.max()
        print x.min()
        print x.mean()
        print
        # '''
        preds = probs.data.cpu().numpy() > 0
        labels = labels.data.cpu().numpy()
        for pred, label in zip(preds, labels):
            pred[label < 0] = -1
            if label.sum() < 0.5:
                continue
            tp = (pred + label == 2).sum()
            tn = (pred + label == 0).sum()
            fp = (pred - label == 1).sum()
            fn = (pred - label == -1).sum()
            precision = 1.0 * tp / (tp + fp + 10e-20)
            recall = 1.0 * tp / (tp + fn + 10e-20)
            f1score = 2. * precision * recall / (precision + recall + 10e-20)
            precision_list.append(precision)
            recall_list.append(recall)
            f1score_list.append(f1score)

        if 'train' in phase and i % 50 == 0:
            loss_mean = np.array(loss_list).mean(0)
            print('step[{:d}] loss: {:3.4f}    pos loss: {:3.4f}   neg loss: {:3.4f}'.format(i, loss_mean[0],
                                                                                             loss_mean[1],
                                                                                             loss_mean[2]))
        # 保存中间结果到 data/middle_result，用于分析
        if i == 0:
            images = images.data.cpu().numpy() * 128 + 128
            if phase == 'pretrain':
                bbox_labels = bbox_labels.data.cpu().numpy()
                seg_labels = seg_labels.data.cpu().numpy()
                seg_output = seg_output.data.cpu().numpy()
            for ii in range(len(images)):
                middle_dir = os.path.join(args.save_dir, 'middle_result')
                if not os.path.exists(middle_dir):
                    os.mkdir(middle_dir)
                middle_dir = os.path.join(middle_dir, phase)
                if not os.path.exists(middle_dir):
                    os.mkdir(middle_dir)
                Image.fromarray(images[ii].astype(np.uint8).transpose(1, 2, 0)).save(
                    os.path.join(middle_dir, str(ii) + '.image.png'))
                if phase == 'pretrain':
                    segi = seg_labels[ii]
                    _segi = np.array([segi, segi, segi]) * 255
                    segi = np.zeros([3, _segi.shape[1] * 2, _segi.shape[2] * 2])
                    for si in range(segi.shape[1]):
                        for sj in range(segi.shape[2]):
                            segi[:, si, sj] = _segi[:, si / 2, sj / 2]
                    Image.fromarray(segi.transpose(1, 2, 0).astype(np.uint8)).save(
                        os.path.join(middle_dir, str(ii) + '.seg.png'))
                    segi = seg_output[ii]
                    _segi = np.array([segi, segi, segi]) * 255
                    segi = np.zeros([3, _segi.shape[1] * 2, _segi.shape[2] * 2])
                    for si in range(segi.shape[1]):
                        for sj in range(segi.shape[2]):
                            segi[:, si, sj] = _segi[:, si / 2, sj / 2]
                    Image.fromarray(segi.transpose(1, 2, 0).astype(np.uint8)).save(
                        os.path.join(middle_dir, str(ii) + '.seg.out.png'))

    f1score = np.mean(f1score_list)
    print 'f1score', f1score
    print 'recall', np.mean(recall_list)
    print 'precision', np.mean(precision_list)
    if 'train' in phase:
        loss_mean = np.array(loss_list).mean(0)
        print 'loss: {:3.4f}    pos loss: {:3.4f}   neg loss: {:3.4f}'.format(loss_mean[0], loss_mean[1], loss_mean[2])

    # 保存模型
    if ('eval' in phase or 'pretrain' in phase) and best_f1score < 2:
        if args.small:
            save_dir = os.path.join(args.save_dir, 'models-small')
        else:
            save_dir = os.path.join(args.save_dir, 'models')
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        if epoch % 3 == 0:
            save_model(save_dir, phase, str(epoch), epoch, f1score, model)
        if f1score > best_f1score:
            save_model(save_dir, phase, 'best_f1score', epoch, f1score, model)
        if args.model == 'resnet':
            tpnr = loss[11] + loss[12]
            # 这里用 best_f1score 也当tpnr好了，懒得改
            if tpnr > best_f1score:
                best_f1score = tpnr
                save_model(save_dir, phase, 'best_tpnr', epoch, f1score, model)
            print 'best tpnr', best_f1score
        else:
            best_f1score = max(best_f1score, f1score)
            if best_f1score < 1:
                print '\n\t{:s}\tbest f1score {:3.4f}\n'.format(phase, best_f1score)
        return best_f1score


def main():
    word_index_dict = json.load(open(args.word_index_json, 'r'))
    num_classes = len(word_index_dict)
    image_label_dict = json.load(open(args.image_label_json, 'r'))
    cudnn.benchmark = True
    if args.model == 'densenet':
        # 两千多种字符，multi-label分类
        model = DenseNet121(num_classes).cuda()
    elif args.model == 'resnet':
        # resnet主要用于文字区域的segmentation以及object detection操作
        model = resnet.ResNet(num_classes=num_classes, args=args).cuda()
    else:
        return
    ##优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    # model = torch.nn.DataParallel(model).cuda()
    loss = Loss().cuda()

    if args.resume:
        state_dict = torch.load(args.resume)
        model.load_state_dict(state_dict['state_dict'])
        best_f1score = state_dict['f1score']
        start_epoch = state_dict['epoch'] + 1
    else:
        best_f1score = 0
        if args.model == 'resnet':
            start_epoch = 100
        else:
            start_epoch = 1
    print 'best_f1score', best_f1score

    # 划分数据集
    # test_filelist1 = sorted(glob(os.path.join('/media/scut214/file/HLG/Chinese_Recognition', 'crop_test1', '*')))
    #data set for single word
    # test_filelist1 = sorted(glob(os.path.join('/media/scut214/file/HLG/Chinese_Recognition', 'crop_test1', '*')))
    test_filelist1 = sorted(glob(os.path.join('../../../intermediate_file', 'images_to_recognition', '*')))


    # test_filelist=os.listdir('/Net/competition/OCR/ocr_densenet/data/dataset/crop')
    # trainval_filelist = sorted(glob(os.path.join(args.data_dir, 'train', '*')))

    # 两种输入size训练
    # 修改为自身数据集的尺寸。
    # train_filelist1: 长宽比小于8:1的图片，经过padding后变成 64*512 的输入
    # train_filelist2: 长宽比大于8:1的图片，经过padding,crop后变成 64*1024的输入
    train_filelist1, train_filelist2 = [], []
    print len(test_filelist1)

    # 黑名单，这些图片的label是有问题的
    black_list = set(json.load(open(args.black_json, 'r'))['black_list'])
    image_hw_ratio_dict = json.load(open(args.image_hw_ratio_json, 'r'))
    test_filelist=[]
    for f in test_filelist1:
        image = f.split('/')[-1]
        if image in black_list:
            continue
        # r = image_hw_ratio_dict[image]
        # if r == 1 or r==2:
        test_filelist.append(f)
            # print r
        # else:
        #     train_filelist2.append(f)
    train_val_filelist = train_filelist1 + train_filelist2
    val_filelist = train_filelist1[-2048:]
    train_filelist1 = train_filelist1[:-2048]
    train_filelist2 = train_filelist2
    image_size = [1300, 64]
    print len(test_filelist)
    if args.phase in ['test', 'val', 'train_val']:
        # 测试输出文字检测结果
        test_dataset = dataloader.DataSet(
            test_filelist,
            image_label_dict,
            num_classes,
            # transform=train_transform,
            args=args,
            image_size=image_size,
            phase='test')
        test_loader = DataLoader(
            dataset=test_dataset,
            batch_size=1,
            shuffle=False,
            num_workers=4,
            pin_memory=True)
        # train_filelist = train_filelist1[-2048:]
        # train_dataset = dataloader.DataSet(
        #     train_filelist,
        #     image_label_dict,
        #     num_classes,
        #     image_size=image_size,
        #     args=args,
        #     phase='test')
        # train_loader = DataLoader(
        #     dataset=train_dataset,
        #     batch_size=1,
        #     shuffle=False,
        #     num_workers=8,
        #     pin_memory=True)
        
        # val_dataset = dataloader.DataSet(
        #     val_filelist,
        #     image_label_dict,
        #     num_classes,
        #     image_size=image_size,
        #     args=args,
        #     phase='test')
        # val_loader = DataLoader(
        #     dataset=val_dataset,
        #     batch_size=1,
        #     shuffle=False,
        #     num_workers=8,
        #     pin_memory=True)
        
        # train_val_dataset = dataloader.DataSet(
        #     train_val_filelist,
        #     image_label_dict,
        #     num_classes,
        #     image_size=image_size,
        #     args=args,
        #     phase='test')
        # train_val_loader = DataLoader(
        #     dataset=train_val_dataset,
        #     batch_size=1,
        #     shuffle=False,
        #     num_workers=8,
        #     pin_memory=True)

        if args.phase == 'test':
            #            test(start_epoch - 1, model, val_loader, 'val')
            test(model, test_loader, 'test')
            # test(start_epoch - 1, model, train_val_loader, 'train_val')
        # elif args.phase == 'val':
        #     test(start_epoch - 1, model, train_loader, 'train')#valid set
        #     test(start_epoch - 1, model, val_loader, 'val')#test set
        # elif args.phase == 'train_val':
        #     test(start_epoch - 1, model, train_val_loader, 'train_val')
        return

    elif args.phase == 'train':

        train_dataset1 = dataloader.DataSet(
            train_filelist1,
            image_label_dict,
            num_classes,
            image_size=image_size,
            args=args,
            phase='train')
        train_loader1 = DataLoader(
            dataset=train_dataset1,
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True)
        #         train_dataset2 = dataloader.DataSet(
        #                 train_filelist2,
        #                 image_label_dict,
        #                 num_classes,
        #                 image_size=(1024,64),
        #                 args=args,
        #                 phase='train')
        #         train_loader2 = DataLoader(
        #                 dataset=train_dataset2,
        #                 batch_size=args.batch_size / 2,
        #                 shuffle=True,
        #                 num_workers=8,
        #                 pin_memory=True)
        val_dataset = dataloader.DataSet(
            val_filelist,
            image_label_dict,
            num_classes,
            image_size=image_size,
            args=args,
            phase='val')
        val_loader = DataLoader(
            dataset=val_dataset,
            batch_size=min(8, args.batch_size),
            shuffle=False,
            num_workers=8,
            pin_memory=True)
        best_f1score = 0
        # eval_mode = 'pretrain-2'
        eval_mode = 'eval'
        for a in range(start_epoch, args.epochs):

            args.epoch = epoch

            if eval_mode == 'eval':
                if best_f1score > 0.8:
                    args.lr = 0.0001
                if best_f1score > 0.7:
                    args.hard_mining = 1

            for param_group in optimizer.param_groups:
                param_group['lr'] = args.lr

            train_eval(epoch, model, train_loader1, loss, optimizer, 2., 'train-1')

            best_f1score = train_eval(epoch, model, val_loader, loss, optimizer, best_f1score,
                                      'eval-{:d}-{:d}'.format(args.batch_size, args.hard_mining))
            print 'best_f1score:',best_f1score
            continue


if __name__ == '__main__':
    print 'eval-{:d}-{:d}'.format(args.batch_size, args.hard_mining)
    main()
