整合去雾网页工具

2026-06-10 17:42:11 +08:00
commit 6db15ebc3f
101 changed files with 10167 additions and 0 deletions
--- a/GCANet/GCANet_train/train.py
+++ b/GCANet/GCANet_train/train.py
@@ -0,0 +1,228 @@
+import os
+import datetime
+import argparse
+import numpy as np
+
+import torch
+import torch.optim as optim
+from torch.autograd import Variable
+from torch.utils.data import DataLoader
+
+from ImagePairPrefixFolder import ImagePairPrefixFolder, var_custom_collate
+from utils import MovingAvg
+from tf_visualizer import TFVisualizer
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--network', default='GCANet')
+parser.add_argument('--name', default='default_exp')
+parser.add_argument('--gpu_ids', default='0')
+parser.add_argument('--epochs', type=int, default=100)
+parser.add_argument('--lr', type=float, default=0.001)
+parser.add_argument('--lr_step', type=int, default=40)
+parser.add_argument('--lr_gamma', type=float, default=0.1)
+parser.add_argument('--weight_decay', type=float, default=0.0005)
+parser.add_argument('--checkpoints_dir', default='checkpoint')
+parser.add_argument('--logDir', default='tblogdir')
+parser.add_argument('--resume_dir', default='')
+parser.add_argument('--resume_epoch', type=int, default=0)
+parser.add_argument('--save_epoch', type=int, default=5)
+parser.add_argument('--save_latest_freq', type=int, default=5000)
+parser.add_argument('--test_epoch', type=int, default=5)
+parser.add_argument('--test_max_size', type=int, default=1080)
+parser.add_argument('--size_unit', type=int,  default=8)
+parser.add_argument('--print_iter', type=int,  default=100)
+parser.add_argument('--input_folder', default='')
+parser.add_argument('--gt_folder', default='')
+parser.add_argument('--test_input_folder', default='')
+parser.add_argument('--test_gt_folder', default='')
+parser.add_argument('--num_workers', type=int, default=16)
+parser.add_argument('--batch_size', type=int, default=4)
+parser.add_argument('--only_residual', action='store_true', help='regress residual rather than image')
+parser.add_argument('--loss_func', default='l2', help='l2|l1')
+parser.add_argument('--inc', type=int, default=3)
+parser.add_argument('--outc', type=int, default=3)
+parser.add_argument('--force_rgb', action='store_true')
+parser.add_argument('--no_edge', action='store_true')
+
+opt = parser.parse_args()
+
+opt.input_folder = os.path.expanduser(opt.input_folder)
+opt.gt_folder = os.path.expanduser(opt.gt_folder)
+opt.test_input_folder = os.path.expanduser(opt.test_input_folder)
+opt.test_gt_folder = os.path.expanduser(opt.test_gt_folder)
+
+
+if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.name)):
+    os.makedirs(os.path.join(opt.checkpoints_dir, opt.name))
+opt.resume_dir = opt.resume_dir if opt.resume_dir != '' else os.path.join(opt.checkpoints_dir, opt.name)
+
+visualizer = TFVisualizer(opt)
+### Log out
+with open(os.path.realpath(__file__), 'r') as fid:
+    visualizer.print_logs(fid.read())
+
+## print argument
+for key, val in vars(opt).items():
+    visualizer.print_logs('%s: %s' % (key, val))
+
+opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(',')]
+assert all(0 <= x <= torch.cuda.device_count() for x in opt.gpu_ids), 'gpu id should ' \
+                                                      'be 0~{0}'.format(torch.cuda.device_count())
+torch.cuda.set_device(opt.gpu_ids[0])
+
+
+train_dataset = ImagePairPrefixFolder(opt.input_folder, opt.gt_folder, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
+train_dataloader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True,
+                              collate_fn=var_custom_collate, pin_memory=True,
+                              num_workers=opt.num_workers)
+
+opt.do_test = opt.test_gt_folder != ''
+if opt.do_test:
+    test_dataset = ImagePairPrefixFolder(opt.test_input_folder, opt.test_gt_folder,
+                                         max_img_size=opt.test_max_size, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
+    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False,
+                                 collate_fn=var_custom_collate, pin_memory=True,
+                                 num_workers=1)
+
+total_inc = opt.inc if opt.no_edge else opt.inc + 1
+if opt.network == 'GCANet':
+    from GCANet import GCANet
+    net = GCANet(in_c=total_inc, out_c=3, only_residual=opt.only_residual)
+else:
+    print('network structure %s not supported' % opt.network)
+    raise ValueError
+
+if opt.loss_func == 'l2':
+    loss_crit = torch.nn.MSELoss()
+elif opt.loss_func == 'l1':
+    loss_crit = torch.nn.SmoothL1Loss()
+else:
+    print('loss_func %s not supported' % opt.loss_func)
+    raise ValueError
+pnsr_crit = torch.nn.MSELoss()
+
+if len(opt.gpu_ids) > 0:
+    net.cuda()
+    if len(opt.gpu_ids) > 1:
+        net = torch.nn.DataParallel(net)
+    loss_crit = loss_crit.cuda()
+    pnsr_crit = pnsr_crit.cuda()
+
+optimizer = optim.Adam(net.parameters(), lr=opt.lr)
+step_optim_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lr_step, gamma=opt.lr_gamma)
+loss_avg = MovingAvg(pool_size=50)
+
+start_epoch = 0
+total_iter = 0
+
+if os.path.exists(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')):
+    print('resuming from latest.pth')
+    latest_info = torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
+    start_epoch = latest_info['epoch']
+    total_iter = latest_info['total_iter']
+    if isinstance(net, torch.nn.DataParallel):
+        net.module.load_state_dict(latest_info['net_state'])
+    else:
+        net.load_state_dict(latest_info['net_state'])
+    optimizer.load_state_dict(latest_info['optim_state'])
+
+if opt.resume_epoch > 0:
+    start_epoch = opt.resume_epoch
+    total_iter = opt.resume_epoch * len(train_dataloader)
+    resume_path = os.path.join(opt.resume_epoch, 'net_epoch_%d.pth') % opt.resume_epoch
+    print('resume from : %s' % resume_path)
+    assert os.path.exists(resume_path), 'cannot find the resume model: %s ' % resume_path
+    if isinstance(net, torch.nn.DataParallel):
+        net.module.load_state_dict(torch.load(resume_path))
+    else:
+        net.load_state_dict(torch.load(resume_path))
+
+for epoch in range(start_epoch, opt.epochs):
+    visualizer.print_logs("Start to train epoch %d" % epoch)
+    net.train()
+    for iter, data in enumerate(train_dataloader):
+        total_iter += 1
+        optimizer.zero_grad()
+        step_optim_scheduler.step(epoch)
+
+        batch_input_img, batch_input_edge,  batch_gt = data
+        if len(opt.gpu_ids) > 0:
+            batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
+
+        if opt.no_edge:
+            batch_input = batch_input_img
+        else:
+            batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
+        batch_input_v = Variable(batch_input)
+        if opt.only_residual:
+            batch_gt_v = Variable(batch_gt - (batch_input_img+128))
+        else:
+            batch_gt_v = Variable(batch_gt)
+
+        pred = net(batch_input_v)
+
+        loss = loss_crit(pred, batch_gt_v)
+        avg_loss = loss_avg.set_curr_val(loss.data)
+
+        loss.backward()
+        optimizer.step()
+
+        if iter % opt.print_iter == 0:
+            visualizer.plot_current_losses(total_iter, { 'loss': loss})
+            visualizer.print_logs('%s Step[%d/%d], lr: %f, mv_avg_loss: %f, loss: %f' %
+                                    (str(datetime.datetime.now()).split(' ')[1], iter, len(train_dataloader),
+                                     step_optim_scheduler.get_lr()[0], avg_loss, loss))
+
+        if total_iter % opt.save_latest_freq == 0:
+            latest_info = {'total_iter': total_iter,
+                           'epoch': epoch,
+                           'optim_state': optimizer.state_dict()}
+            if len(opt.gpu_ids) > 1:
+                latest_info['net_state'] = net.module.state_dict()
+            else:
+                latest_info['net_state'] = net.state_dict()
+            print('save lastest model.')
+            torch.save(latest_info, os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
+
+    if (epoch+1) % opt.save_epoch == 0 :
+        visualizer.print_logs('saving model for epoch %d' % epoch)
+        if len(opt.gpu_ids) > 1:
+            torch.save(net.module.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch+1)))
+        else:
+            torch.save(net.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch + 1)))
+
+    if opt.do_test:        
+        avg_psnr = 0
+        task_cnt = 0
+        net.eval()
+        with torch.no_grad():
+            for iter, data in enumerate(test_dataloader):
+                batch_input_img, batch_input_edge,  batch_gt = data
+                if len(opt.gpu_ids) > 0:
+                    batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
+
+                if opt.no_edge:
+                    batch_input = batch_input_img
+                else:
+                    batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
+                batch_input_v = Variable(batch_input)
+                batch_gt_v = Variable(batch_gt)
+
+
+                pred = net(batch_input_v)
+
+                if opt.only_residual:
+                    loss = pnsr_crit(pred+Variable(batch_input_img+128), batch_gt_v)
+                else:
+                    loss = pnsr_crit(pred, batch_gt_v)
+                avg_psnr += 10 * np.log10(255 * 255 / loss.item())
+                task_cnt += 1
+
+        visualizer.print_logs('Testing for epoch: %d' % epoch)
+        visualizer.print_logs('Average test PNSR is %f for %d images' % (avg_psnr/task_cnt, task_cnt))
+
+
+
+
+
+