import sys
from table_line import model, dice_coef_loss, dice_coef
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint,ReduceLROnPlateau
from sklearn.model_selection import train_test_split
from glob import glob
from image import gen, gen_origin

if __name__=='__main__':
    filepath = 'models/table-line.h5'##模型权重存放位置

    checkpointer = ModelCheckpoint(filepath=filepath,monitor='loss',verbose=0,save_weights_only=True, save_best_only=True)
    rlu = ReduceLROnPlateau(monitor='loss', factor=0.1, patience=5, verbose=0, mode='auto', cooldown=0, min_lr=0)
    model.compile(optimizer=Adam(lr=0.0003), loss=dice_coef_loss(), metrics=['acc', dice_coef])

    paths = glob('./train/dataset-line/6/*.json') ##table line dataset label with labelme
    trainP,testP = train_test_split(paths,test_size=0.1)
    print('total:',len(paths),'train:',len(trainP),'test:',len(testP))
    batchsize=4
    trainloader = gen_origin(trainP,batchsize=batchsize,linetype=1)
    testloader = gen_origin(testP,batchsize=batchsize,linetype=1)
    model.fit_generator(trainloader,
                        steps_per_epoch=max(1,len(trainP)//batchsize),
                        callbacks=[checkpointer],
                        validation_data=testloader,
                        validation_steps=max(1,len(testP)//batchsize),
                        epochs=30)