def main():
#### setup options of three networks
parser = argparse.ArgumentParser()
parser.add_argument('-opt_P', type=str, help='Path to option YMAL file of Predictor.')
parser.add_argument('-opt_C', type=str, help='Path to option YMAL file of Corrector.')
parser.add_argument('-opt_F', type=str, help='Path to option YMAL file of SFTMD_Net.')
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt_P = option.parse(args.opt_P, is_train=True)
opt_C = option.parse(args.opt_C, is_train=True)
opt_F = option.parse(args.opt_F, is_train=True)
# convert to NoneDict, which returns None for missing keys
opt_P = option.dict_to_nonedict(opt_P)
opt_C = option.dict_to_nonedict(opt_C)
opt_F = option.dict_to_nonedict(opt_F)
# choose small opt for SFTMD test
opt_F = opt_F['sftmd']
#### random seed
seed = opt_P['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
util.set_random_seed(seed)
# create PCA matrix of enough kernel
batch_ker = util.random_batch_kernel(batch=30000, l=opt_P['kernel_size'], sig_min=0.2, sig_max=4.0, rate_iso=1.0, scaling=3, tensor=False)
print('batch kernel shape: {}'.format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
pca_matrix = util.PCA(batch_ker, k=opt_P['code_length']).float()
print('PCA matrix shape: {}'.format(pca_matrix.shape))
#### distributed training settings
if args.launcher == 'none': # disabled distributed training
opt_P['dist'] = False
opt_F['dist'] = False
opt_C['dist'] = False
rank = -1
print('Disabled distributed training.')
else:
opt_P['dist'] = True
opt_F['dist'] = True
opt_C['dist'] = True
init_dist()
world_size = torch.distributed.get_world_size() #Returns the number of processes in the current process group
rank = torch.distributed.get_rank() #Returns the rank of current process group
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
###### SFTMD train ######
SFTMD_train(opt_F, rank, world_size, pca_matrix)
###### Predictor&Corrector train ######
IKC_train(opt_P, opt_C, opt_F, rank, world_size, pca_matrix)
def main():
#### setup options of three networks
parser = argparse.ArgumentParser()
parser.add_argument('-opt_P',
type=str,
help='Path to option YMAL file of Predictor.')
parser.add_argument('-opt_C',
type=str,
help='Path to option YMAL file of Corrector.')
parser.add_argument('-opt_F',
type=str,
help='Path to option YMAL file of SFTMD_Net.')
parser.add_argument('--launcher',
choices=['none', 'pytorch'],
default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt_P = option.parse(args.opt_P, is_train=True)
opt_C = option.parse(args.opt_C, is_train=True)
opt_F = option.parse(args.opt_F, is_train=True)
# convert to NoneDict, which returns None for missing keys
opt_P = option.dict_to_nonedict(opt_P)
opt_C = option.dict_to_nonedict(opt_C)
opt_F = option.dict_to_nonedict(opt_F)
# choose small opt for SFTMD test, fill path of pre-trained model_F
opt_F = opt_F['sftmd']
# create PCA matrix of enough kernel
batch_ker = util.random_batch_kernel(batch=30000,
l=opt_P['kernel_size'],
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
tensor=False)
print('batch kernel shape: {}'.format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
pca_matrix = util.PCA(batch_ker, k=opt_P['code_length']).float()
print('PCA matrix shape: {}'.format(pca_matrix.shape))
#### distributed training settings
if args.launcher == 'none': # disabled distributed training
opt_P['dist'] = False
opt_F['dist'] = False
opt_C['dist'] = False
rank = -1
print('Disabled distributed training.')
else:
opt_P['dist'] = True
opt_F['dist'] = True
opt_C['dist'] = True
init_dist()
world_size = torch.distributed.get_world_size(
) #Returns the number of processes in the current process group
rank = torch.distributed.get_rank(
) #Returns the rank of current process group
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
###### Predictor&Corrector train ######
#### loading resume state if exists
if opt_P['path'].get('resume_state', None):
# distributed resuming: all load into default GPU
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt_P['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
option.check_resume(opt_P,
resume_state['iter']) # check resume options
else:
resume_state = None
#### mkdir and loggers
if rank <= 0: # normal training (rank -1) OR distributed training (rank 0-7)
if resume_state is None:
# Predictor path
util.mkdir_and_rename(
opt_P['path']
['experiments_root']) # rename experiment folder if exists
util.mkdirs(
(path for key, path in opt_P['path'].items()
if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key))
# Corrector path
util.mkdir_and_rename(
opt_C['path']
['experiments_root']) # rename experiment folder if exists
util.mkdirs(
(path for key, path in opt_C['path'].items()
if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key))
# config loggers. Before it, the log will not work
util.setup_logger('base',
opt_P['path']['log'],
'train_' + opt_P['name'],
level=logging.INFO,
screen=True,
tofile=True)
util.setup_logger('val',
opt_P['path']['log'],
'val_' + opt_P['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_P))
logger.info(option.dict2str(opt_C))
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['name']:
version = float(torch.__version__[0:3])
if version >= 1.1: # PyTorch 1.1
from torch.utils.tensorboard import SummaryWriter
else:
logger.info(
'You are using PyTorch {}. Tensorboard will use [tensorboardX]'
.format(version))
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt_P['name'])
else:
util.setup_logger('base',
opt_P['path']['log'],
'train',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
#### random seed
seed = opt_P['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
if rank <= 0:
logger.info('Random seed: {}'.format(seed))
util.set_random_seed(seed)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
#### create train and val dataloader
dataset_ratio = 200 # enlarge the size of each epoch
for phase, dataset_opt in opt_P['datasets'].items():
if phase == 'train':
train_set = create_dataset(dataset_opt)
train_size = int(
math.ceil(len(train_set) / dataset_opt['batch_size']))
total_iters = int(opt_P['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
if opt_P['dist']:
train_sampler = DistIterSampler(train_set, world_size, rank,
dataset_ratio)
total_epochs = int(
math.ceil(total_iters / (train_size * dataset_ratio)))
else:
train_sampler = None
train_loader = create_dataloader(train_set, dataset_opt, opt_P,
train_sampler)
if rank <= 0:
logger.info(
'Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, total_iters))
elif phase == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt, opt_P, None)
if rank <= 0:
logger.info('Number of val images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(val_set)))
else:
raise NotImplementedError(
'Phase [{:s}] is not recognized.'.format(phase))
assert train_loader is not None
assert val_loader is not None
#### create model
model_F = create_model(opt_F) #load pretrained model of SFTMD
model_P = create_model(opt_P)
model_C = create_model(opt_C)
#### resume training
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model_P.resume_training(
resume_state) # handle optimizers and schedulers
else:
current_step = 0
start_epoch = 0
#### training
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
for epoch in range(start_epoch, total_epochs + 1):
if opt_P['dist']:
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
#### update learning rate, schedulers
# model.update_learning_rate(current_step, warmup_iter=opt_P['train']['warmup_iter'])
#### preprocessing for LR_img and kernel map
prepro = util.SRMDPreprocessing(opt_P['scale'],
pca_matrix,
para_input=opt_P['code_length'],
kernel=opt_P['kernel_size'],
noise=False,
cuda=True,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
rate_cln=0.2,
noise_high=0.0)
LR_img, ker_map = prepro(train_data['GT'])
#### training Predictor
model_P.feed_data(LR_img, ker_map)
model_P.optimize_parameters(current_step)
P_visuals = model_P.get_current_visuals()
est_ker_map = P_visuals['Batch_est_ker_map']
#### log of model_P
if current_step % opt_P['logger']['print_freq'] == 0:
logs = model_P.get_current_log()
message = 'Predictor <epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model_P.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
#### training Corrector
for step in range(opt_C['step']):
# test SFTMD for corresponding SR image
model_F.feed_data(train_data, LR_img, est_ker_map)
model_F.test()
F_visuals = model_F.get_current_visuals()
SR_img = F_visuals['Batch_SR']
# Test SFTMD to produce SR images
# train corrector given SR image and estimated kernel map
model_C.feed_data(SR_img, est_ker_map, ker_map)
model_C.optimize_parameters(current_step)
C_visuals = model_C.get_current_visuals()
est_ker_map = C_visuals['Batch_est_ker_map']
#### log of model_C
if current_step % opt_C['logger']['print_freq'] == 0:
logs = model_C.get_current_log()
message = 'Corrector <epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step,
model_C.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt_C['use_tb_logger'] and 'debug' not in opt_C[
'name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# validation, to produce ker_map_list(fake)
if current_step % opt_P['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for _, val_data in enumerate(val_loader):
prepro = util.SRMDPreprocessing(
opt_P['scale'],
pca_matrix,
para_input=opt_P['code_length'],
kernel=opt_P['kernel_size'],
noise=False,
cuda=True,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
rate_cln=0.2,
noise_high=0.0)
LR_img, ker_map = prepro(val_data['GT'])
single_img_psnr = 0.0
# valid Predictor
model_P.feed_data(LR_img, ker_map)
model_P.test()
P_visuals = model_P.get_current_visuals()
est_ker_map = P_visuals['Batch_est_ker_map']
for step in range(opt_C['step']):
step += 1
idx += 1
model_F.feed_data(val_data, LR_img, est_ker_map)
model_F.test()
F_visuals = model_F.get_current_visuals()
SR_img = F_visuals['Batch_SR']
# Test SFTMD to produce SR images
model_C.feed_data(SR_img, est_ker_map, ker_map)
model_C.test()
C_visuals = model_C.get_current_visuals()
est_ker_map = C_visuals['Batch_est_ker_map']
sr_img = util.tensor2img(F_visuals['SR']) # uint8
gt_img = util.tensor2img(F_visuals['GT']) # uint8
# Save SR images for reference
img_name = os.path.splitext(
os.path.basename(val_data['LQ_path'][0]))[0]
img_dir = os.path.join(opt_P['path']['val_images'],
img_name)
# img_dir = os.path.join(opt_F['path']['val_images'], str(current_step), '_', str(step))
util.mkdir(img_dir)
save_img_path = os.path.join(
img_dir, '{:s}_{:d}_{:d}.png'.format(
img_name, current_step, step))
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt_P['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size, :]
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size, :]
step_psnr = util.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255)
logger.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> img:{:s}, psnr: {:.4f}'
.format(epoch, current_step, step, img_name,
step_psnr))
single_img_psnr += step_psnr
avg_psnr += util.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255)
avg_signle_img_psnr = single_img_psnr / step
logger.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> img:{:s}, average psnr: {:.4f}'
.format(epoch, current_step, step, img_name,
avg_signle_img_psnr))
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4f}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> psnr: {:.4f}'.
format(epoch, current_step, step, avg_psnr))
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
#### save models and training states
if current_step % opt_P['logger']['save_checkpoint_freq'] == 0:
if rank <= 0:
logger.info('Saving models and training states.')
model_P.save(current_step)
model_P.save_training_state(epoch, current_step)
model_C.save(current_step)
model_C.save_training_state(epoch, current_step)
if rank <= 0:
logger.info('Saving the final model.')
model_P.save('latest')
model_C.save('latest')
logger.info('End of Predictor and Corrector training.')
tb_logger.close()
tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_F))
# create PCA matrix of enough kernel
batch_ker = util.random_batch_kernel(batch=30000,
l=21,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
tensor=False)
print('batch kernel shape: {}'.format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
pca_matrix = util.PCA(batch_ker, k=10).float()
print('PCA matrix shape: {}'.format(pca_matrix.shape))
#### Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt_F['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
logger.info('Number of test images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
# load pretrained model by default
model_F = create_model(opt_F)
for test_loader in test_loaders:
parser = argparse.ArgumentParser()
parser.add_argument("-opt",
type=str,
help="Path to option YMAL file of SFTMD_Net.")
parser.add_argument("--launcher",
choices=["none", "pytorch"],
default="none",
help="job launcher")
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()
opt = option.parse(args.opt, is_train=True)
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
batch_ker = util.random_batch_kernel(batch=30000,
l=11,
sig_min=0.6,
sig_max=5,
rate_iso=0,
scaling=3,
tensor=False)
print("batch kernel shape: {}".format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
pca_matrix = util.PCA(batch_ker, k=opt["code_length"]).float()
print("PCA matrix shape: {}".format(pca_matrix.shape))
torch.save(pca_matrix, "./pca_matrix.pth")
print("Save PCA matrix at: ./pca_matrix.pth")
def degradation_model():
# set parameters
scale = 4
mod_scale = 4
# set data dir
sourcedir = 'D:/NEU/ImageRestoration/div2k/Set5/HR/'
savedir = 'D:/NEU/ImageRestoration/div2k/Set5/'
# set random seed和
util.set_random_seed(0)
# load PCA matrix of enough kernelh
# print('load PCA matrix')
# pca_matrix = torch.load('/media/sdc/yjchai/IKC/codes/pca_matrix.pth', map_location=lambda storage, loc: storage)
# print('PCA matrix shape: {}'.format(pca_matrix.shape))
# 初始化kernelmap
batch_ker = util.random_batch_kernel(batch=1000,
k=15,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
tensor=False)
print('batch kernel shape: {}'.format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
dim_pca = 15
# torch.Size([225, 15])
pca_matrix = util.PCA(batch_ker, dim_pca).float()
print('PCA matrix shape: {}'.format(pca_matrix.shape))
# saveHRpath = os.path.join(savedir, 'HR')
# saveLRpath = os.path.join(savedir, 'LR', 'x' + str(scale))
# saveBicpath = os.path.join(savedir, 'Bic', 'x' + str(scale))
saveLRblurpath = os.path.join(savedir, 'LRblur_sig0.6', 'X' + str(scale))
if not os.path.isdir(sourcedir):
print('Error: No source data found')
exit(0)
if not os.path.isdir(savedir):
os.mkdir(savedir)
# if not os.path.isdir(os.path.join(savedir, 'HR')):
# os.mkdir(os.path.join(savedir, 'HR'))
# if not os.path.isdir(os.path.join(savedir, 'LR')):
# os.mkdir(os.path.join(savedir, 'LR'))
# if not os.path.isdir(os.path.join(savedir, 'Bic')):
# os.mkdir(os.path.join(savedir, 'Bic'))
if not os.path.isdir(os.path.join(savedir, 'LRblur_sig0.6')):
os.mkdir(os.path.join(savedir, 'LRblur_sig0.6'))
# if not os.path.isdir(saveHRpath):
# os.mkdir(saveHRpath)
# else:
# print('It will cover ' + str(saveHRpath))
# if not os.path.isdir(saveLRpath):
# os.mkdir(saveLRpath)
# else:
# print('It will cover ' + str(saveLRpath))
# if not os.path.isdir(saveBicpath):
# os.mkdir(saveBicpath)
# else:
# print('It will cover ' + str(saveBicpath))
if not os.path.isdir(saveLRblurpath):
os.mkdir(saveLRblurpath)
else:
print('It will cover ' + str(saveLRblurpath))
# 记得更改图片格式，不然无法生成
filepaths = sorted(
[f for f in os.listdir(sourcedir) if f.endswith('.bmp')])
num_files = len(filepaths)
# prepare data with augementation
for i in range(num_files):
filename = filepaths[i]
print('No.{} -- Processing {}'.format((i + 1), filename))
# read image
image_HR = cv2.imread(os.path.join(sourcedir, filename))
image = cv2.imread(os.path.join(sourcedir, filename))
width = int(np.floor(image.shape[1] / mod_scale))
height = int(np.floor(image.shape[0] / mod_scale))
# modcrop
if len(image.shape) == 3:
image_HR = image[0:mod_scale * height, 0:mod_scale * width, :]
else:
image_HR = image[0:mod_scale * height, 0:mod_scale * width]
# LR_blur, by random gaussian kernel
img_HR = util.img2tensor(image_HR)
C, H, W = img_HR.size()
# sig_list = [1.8, 2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2]
sig = 0.6
prepro = util.SRMDPreprocessing(
scale,
pca_matrix,
random=False,
kernel=15,
noise=False,
cuda=True,
sig=sig,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=4,
rate_cln=0.2,
noise_high=0.0) #random(sig_min, sig_max) | stable kernel(sig)
LR_img, ker_map = prepro(img_HR.view(1, C, H, W))
image_LR_blur = util.tensor2img(LR_img)
cv2.imwrite(
os.path.join(saveLRblurpath, 'sig{}_'.format(str(sig)) + filename),
image_LR_blur)
# LR
# image_LR = imresize_np(image_HR, 1 / scale, True)
# bic
# image_Bic = imresize_np(image_LR, scale, True)
# cv2.imwrite(os.path.join(saveHRpath, filename), image_HR)
# cv2.imwrite(os.path.join(saveLRpath, filename), image_LR)
# cv2.imwrite(os.path.join(saveBicpath, filename), image_Bic)
# kernel_map_tensor[i] = ker_map
# save dataset corresponding kernel maps
# 1016windows下路径不可含特殊符号，
# torch.save(kernel_map_tensor, 'D:/NEU/ImageRestoration/datasets/try/try_sig2.6_kermap.pth')
print("Image Blurring & Down smaple Done: X" + str(scale))