from keras.callbacks import ModelCheckpoint,Callback,LearningRateScheduler,TensorBoard
from keras.models import load_model
import random
import numpy as np
from scipy import misc
import gc
from keras.optimizers import Adam
from imageio import imread
from datetime import datetime
import os
import json
import models
from utils import DataLoader, LrPolicy
from config import Config
import argparse

def get_parser():
    
    parser = argparse.ArgumentParser('train')
    parser.add_argument('--configPath', '-c', required=True)
    return parser

def train(args=None):
    parser = get_parser()
    args = parser.parse_args(args)
    conf=Config()
    conf.load(args.configPath)
    time=datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    trainString="%s_%s_%s_%s" % (conf.model,conf.optimizer,str(conf.lr),time)
    os.makedirs(conf.logPath+"/"+trainString)
    conf.save(conf.logPath+"/"+trainString+'/config.json')
    print('Compiling model...')
    model_checkpoint = ModelCheckpoint(conf.logPath+"/"+trainString+'/Checkpoint-{epoch:02d}-{val_loss:.2f}.hdf5', monitor='val_loss', save_best_only=False, save_weights_only=True)
    change_lr = LearningRateScheduler(LrPolicy(conf.lr).stepDecay)
    tbCallBack=TensorBoard(log_dir=conf.logPath+"/"+trainString+'/logs', histogram_freq=0,  write_graph=True, write_images=True)
    model=models.modelCreator(conf.model,conf.inputShape,conf.classes,conf.pretrainedModel)
    model.compile(optimizer = conf.optimizer, loss = conf.loss)
    data = [conf.trainDataPath+"/"+f for f in os.listdir(conf.trainDataPath) if '.jpg' in f]
    random.shuffle(data)
    thr=int(len(data)*conf.validationSplit)
    trainData=data[thr:]
    valData=data[:thr]
    trainDataLoader=DataLoader(conf.batchSize,conf.inputShape,trainData,conf.guaMaxValue)
    validationDataLoader=DataLoader(conf.batchSize,conf.inputShape,valData,conf.guaMaxValue)
    print('Fitting model...')
    model.fit_generator(generator=trainDataLoader.generator(),
                    validation_data=validationDataLoader.generator(),
                    steps_per_epoch=len(trainData)//conf.batchSize,
                    validation_steps=len(valData)//conf.batchSize,
                    epochs=conf.epoches,
                    verbose=1,
                    initial_epoch=0,
                    callbacks = [model_checkpoint, change_lr,tbCallBack]
                    )

if __name__ == "__main__":
   train()

[((19, 13, 19, 45), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({(19, 37, 19, 44): '"""train"""'}, {}), "('train')", False, 'import argparse\n'), ((26, 9, 26, 17), 'config.Config', 'Config', ({}, {}), '()', False, 'from config import Config\n'), ((30, 4, 30, 45), 'os.makedirs', 'os.makedirs', ({(30, 16, 30, 44): "(conf.logPath + '/' + trainString)"}, {}), "(conf.logPath + '/' + trainString)", False, 'import os\n'), ((33, 23, 33, 180), 'keras.callbacks.ModelCheckpoint', 'ModelCheckpoint', (), '', False, 'from keras.callbacks import ModelCheckpoint, Callback, LearningRateScheduler, TensorBoard\n'), ((35, 15, 35, 128), 'keras.callbacks.TensorBoard', 'TensorBoard', (), '', False, 'from keras.callbacks import ModelCheckpoint, Callback, LearningRateScheduler, TensorBoard\n'), ((36, 10, 36, 91), 'models.modelCreator', 'models.modelCreator', ({(36, 30, 36, 40): 'conf.model', (36, 41, 36, 56): 'conf.inputShape', (36, 57, 36, 69): 'conf.classes', (36, 70, 36, 90): 'conf.pretrainedModel'}, {}), '(conf.model, conf.inputShape, conf.classes, conf.\n    pretrainedModel)', False, 'import models\n'), ((39, 4, 39, 24), 'random.shuffle', 'random.shuffle', ({(39, 19, 39, 23): 'data'}, {}), '(data)', False, 'import random\n'), ((43, 20, 43, 89), 'utils.DataLoader', 'DataLoader', ({(43, 31, 43, 45): 'conf.batchSize', (43, 46, 43, 61): 'conf.inputShape', (43, 62, 43, 71): 'trainData', (43, 72, 43, 88): 'conf.guaMaxValue'}, {}), '(conf.batchSize, conf.inputShape, trainData, conf.guaMaxValue)', False, 'from utils import DataLoader, LrPolicy\n'), ((44, 25, 44, 92), 'utils.DataLoader', 'DataLoader', ({(44, 36, 44, 50): 'conf.batchSize', (44, 51, 44, 66): 'conf.inputShape', (44, 67, 44, 74): 'valData', (44, 75, 44, 91): 'conf.guaMaxValue'}, {}), '(conf.batchSize, conf.inputShape, valData, conf.guaMaxValue)', False, 'from utils import DataLoader, LrPolicy\n'), ((28, 9, 28, 23), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n'), ((34, 38, 34, 55), 'utils.LrPolicy', 'LrPolicy', ({(34, 47, 34, 54): 'conf.lr'}, {}), '(conf.lr)', False, 'from utils import DataLoader, LrPolicy\n'), ((38, 46, 38, 76), 'os.listdir', 'os.listdir', ({(38, 57, 38, 75): 'conf.trainDataPath'}, {}), '(conf.trainDataPath)', False, 'import os\n')]