import os, time, torch, argparse import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision.utils import save_image as imwrite import numpy as np from torchvision import transforms from makedataset import Dataset from utils.utils import print_args, load_restore_ckpt_with_optim, load_embedder_ckpt, adjust_learning_rate, data_process, tensor_metric, load_excel, save_checkpoint from model.loss import Total_loss from PIL import Image transform_resize = transforms.Compose([ transforms.Resize([224,224]), transforms.ToTensor() ]) def main(args): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print('> Model Initialization...') embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=args.embedder_model_path) restorer, optimizer, cur_epoch = load_restore_ckpt_with_optim(device, freeze_model=False, ckpt_name=args.restore_model_path, lr=args.lr) loss = Total_loss(args) print('> Loading dataset...') data = Dataset(args.train_input) dataset = DataLoader(dataset=data, num_workers=args.num_works, batch_size=args.bs, shuffle=True) print('> Start training...') start_all = time.time() train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device) end_all = time.time() print('Whloe Training Time:' +str(end_all-start_all)+'s.') def train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device): metric = [] for epoch in range(cur_epoch, args.epoch): optimizer = adjust_learning_rate(optimizer, epoch, args.adjust_lr) learnrate = optimizer.param_groups[-1]['lr'] restorer.train() for i, data in enumerate(dataset,0): pos, inp, neg = data_process(data, args, device) text_embedding,_,_ = embedder(inp[1],'text_encoder') out = restorer(inp[0], text_embedding) restorer.zero_grad() total_loss = loss(inp, pos, neg, out) total_loss.backward() optimizer.step() mse = tensor_metric(pos,out, 'MSE', data_range=1) psnr = tensor_metric(pos,out, 'PSNR', data_range=1) ssim = tensor_metric(pos,out, 'SSIM', data_range=1) print("[epoch %d][%d/%d] lr :%f Floss: %.4f MSE: %.4f PSNR: %.4f SSIM: %.4f"%(epoch+1, i+1, \ len(dataset), learnrate, total_loss.item(), mse, psnr, ssim)) psnr_t1, ssim_t1, psnr_t2, ssim_t2 = test(args, restorer, embedder, device, epoch) metric.append([psnr_t1, ssim_t1, psnr_t2, ssim_t2]) print("[epoch %d] Test images PSNR1: %.4f SSIM1: %.4f"%(epoch+1, psnr_t1,ssim_t1)) load_excel(metric) save_checkpoint({'epoch': epoch + 1,'state_dict': restorer.state_dict(),'optimizer' : optimizer.state_dict()},\ args.save_model_path, epoch+1, psnr_t1,ssim_t1,psnr_t2,ssim_t2) def test(args, restorer, embedder, device, epoch=-1): combine_type = args.degr_type psnr_1, psnr_2, ssim_1, ssim_2 = 0, 0, 0, 0 os.makedirs(args.output,exist_ok=True) for i in range(len(combine_type)-1): file_list = os.listdir(f'{args.test_input}/{combine_type[i+1]}/') for j in range(len(file_list)): hq = Image.open(f'{args.test_input}/{combine_type[0]}/{file_list[j]}') lq = Image.open(f'{args.test_input}/{combine_type[i+1]}/{file_list[j]}') restorer.eval() with torch.no_grad(): lq_re = torch.Tensor((np.array(lq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu") lq_em = transform_resize(lq).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu") hq = torch.Tensor((np.array(hq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu") starttime = time.time() text_embedding_1,_,text_1 = embedder([combine_type[i+1]],'text_encoder') text_embedding_2,_, text_2 = embedder(lq_em,'image_encoder') out_1 = restorer(lq_re, text_embedding_1) if text_1 != text_2: print(text_1, text_2) out_2 = restorer(lq_re, text_embedding_2) else: out_2 = out_1 endtime1 = time.time() imwrite(torch.cat((lq_re, out_1, out_2, hq), dim=3), args.output \ + file_list[j][:-4] + '_' + str(epoch) + '_' + combine_type[i+1] + '.png', range=(0, 1)) psnr_1 += tensor_metric(hq, out_1, 'PSNR', data_range=1) ssim_1 += tensor_metric(hq, out_1, 'SSIM', data_range=1) psnr_2 += tensor_metric(hq, out_2, 'PSNR', data_range=1) ssim_2 += tensor_metric(hq, out_2, 'SSIM', data_range=1) print('The ' + file_list[j][:-4] + ' Time:' + str(endtime1 - starttime) + 's.') return psnr_1 / (len(file_list)*len(combine_type)), ssim_1 / (len(file_list)*len(combine_type)),\ psnr_2 / (len(file_list)*len(combine_type)), ssim_2 / (len(file_list)*len(combine_type)) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = "0" if __name__ == '__main__': parser = argparse.ArgumentParser(description = "OneRestore Training") # load model parser.add_argument("--embedder-model-path", type=str, default = "./ckpts/embedder_model.tar", help = 'embedder model path') parser.add_argument("--restore-model-path", type=str, default = None, help = 'restore model path') parser.add_argument("--save-model-path", type=str, default = "./ckpts/", help = 'restore model path') parser.add_argument("--epoch", type=int, default = 300, help = 'epoch number') parser.add_argument("--bs", type=int, default = 4, help = 'batchsize') parser.add_argument("--lr", type=float, default = 1e-4, help = 'learning rate') parser.add_argument("--adjust-lr", type=int, default = 30, help = 'adjust learning rate') parser.add_argument("--num-works", type=int, default = 4, help = 'number works') parser.add_argument("--loss-weight", type=tuple, default = (0.6,0.3,0.1), help = 'loss weights') parser.add_argument("--degr-type", type=list, default = ['clear', 'low', 'haze', 'rain', 'snow',\ 'low_haze', 'low_rain', 'low_snow', 'haze_rain', 'haze_snow', 'low_haze_rain', 'low_haze_snow'], help = 'degradation type') parser.add_argument("--train-input", type=str, default = "./dataset.h5", help = 'train data') parser.add_argument("--test-input", type=str, default = "./data/CDD-11_test", help = 'test path') parser.add_argument("--output", type=str, default = "./result/", help = 'output path') argspar = parser.parse_args() print_args(argspar) main(argspar)