import os from databuilder.rvl_cdip import build_dataset from utils.configloader import ConfigLoader from models.srcnn import Model # from models.srresnet import Model DATA_DIR = '/home/super-videt/Storage/Resource/CV/dataset' NEW_DIR = './datasets' if __name__ == '__main__': # os.environ["CUDA_VISIBLE_DEVICES"] = "0" srcnn_path = "results/checkpoints/srcnn/srcnn_param_batch4_lr0.001_epoch10.pkl" cfg = ConfigLoader('configs/srcnn.cfg', 'default') cfg = ConfigLoader('configs/srcnn.cfg', 'rvl_cdip') model = Model(cfg) model.train(srcnn_path=srcnn_path, random_scale=False, rotate=False, fliplr=False, fliptb=False) model.test(srcnn_path=srcnn_path) # os.environ["CUDA_VISIBLE_DEVICES"] = "1" # srcnn_path = "results/checkpoints/srresnet/srresnet_param_batch4_lr0.001_epoch3.pkl" # cfg = ConfigLoader('configs/srresnet.cfg', 'default') # cfg = ConfigLoader('configs/srresnet.cfg', 'rvl_cdip') # model = Model(cfg) # model.train(srcnn_path=srcnn_path,
if __name__ == '__main__': # build_dataset(data_dir=DATA_DIR, new_dir=NEW_DIR, mode='test') # os.environ["CUDA_VISIBLE_DEVICES"] = "1" edgenet_path = "edgenet_param_batch4_lr0.002_epoch100.pkl" srcnn_path = "results/checkpoints/srcnn/srcnn_param_batch4_lr0.001_epoch3.pkl" sr2x1_path = "results/checkpoints/srunitnet/srnet2x1_param_batch4_lr0.0005_epoch1.pkl" sr2x2_path = "results/checkpoints/srunitnet/srnet2x2_param_batch4_lr0.0005_epoch1.pkl" srresnet_path = "results/checkpoints/srresnet/srresnet_param_batch4_lr0.001_epoch3.pkl" sr2x1_edge_path = "results/checkpoints/srunitnet/srnet2x1_param_batch4_lr0.0005_epoch47.pkl" sr2x2_edge_path = "results/checkpoints/srunitnet/srnet2x2_param_batch4_lr0.0005_epoch47.pkl" # sr2x1_path = None # sr2x2_path = None cfg = ConfigLoader('configs/default.cfg', 'document') cfg = ConfigLoader('configs/default.cfg', 'rvl_cdip') # cfg = ConfigLoader('configs/default.cfg', 'test') model = Model(cfg) model.train(sr2x1_path=sr2x1_path, sr2x2_path=sr2x2_path, edgenetpath=edgenet_path, srcnn_path=srcnn_path, srresnet_path=srresnet_path, random_scale=False, is_fine_tune=False, rotate=False, fliplr=False, fliptb=False) # model.train(sr2x1_path=sr2x1_path, sr2x2_path=sr2x2_path, edgenetpath=edgenet_path, # random_scale=False, is_fine_tune=True, rotate=False, fliplr=False, fliptb=False) # model.train(sr2x1_path=sr2x1_path, sr2x2_path=sr2x2_path, edgenetpath=edgenet_path, # random_scale=False, is_fine_tune=True, rotate=False, fliplr=False, fliptb=False) # model.test(sr2x1_path=sr2x1_path, sr2x2_path=sr2x2_path) # model.test(srcnn_path=srcnn_path) # sr2x1_edge_path = "results/checkpoints/srunitnet/srnet2x1_param_batch4_lr0.0005_epoch47.pkl" # sr2x2_edge_path = "results/checkpoints/srunitnet/srnet2x2_param_batch4_lr0.0005_epoch47.pkl" # sr2x1_none_path = "results/checkpoints/srunitnet/srnet2x1_param_batch4_lr0.001_epoch25.pkl"