123456789101112131415161718192021222324252627 |
- 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)
|