def evaluate(model, dataloader, epoch, writer, logger, data_name='val'):
save_root = os.path.join(opt.result_dir, opt.tag, str(epoch), data_name)
utils.try_make_dir(save_root)
total_psnr = 0.0
total_ssim = 0.0
ct_num = 0
# print('Start testing ' + tag + '...')
for i, sample in enumerate(dataloader):
utils.progress_bar(i, len(dataloader), 'Eva... ')
path = sample['path']
with torch.no_grad():
recovered = model(sample)
if data_name == 'val':
label = sample['label']
label = tensor2im(label)
recovered = tensor2im(recovered)
ct_num += 1
total_psnr += psnr(recovered, label, data_range=255)
total_ssim += ski_ssim(recovered,
label,
data_range=255,
multichannel=True)
save_dst = os.path.join(save_root,
utils.get_file_name(path[0]) + '.png')
Image.fromarray(recovered).save(save_dst)
elif data_name == 'test':
pass
else:
raise Exception('Unknown dataset name: %s.' % data_name)
# 保存结果
save_dst = os.path.join(save_root,
utils.get_file_name(path[0]) + '.png')
Image.fromarray(recovered).save(save_dst)
if data_name == 'val':
ave_psnr = total_psnr / float(ct_num)
ave_ssim = total_ssim / float(ct_num)
# write_loss(writer, f'val/{data_name}', 'psnr', total_psnr / float(ct_num), epochs)
logger.info(f'Eva({data_name}) epoch {epoch}, psnr: {ave_psnr}.')
logger.info(f'Eva({data_name}) epoch {epoch}, ssim: {ave_ssim}.')
return f'{ave_ssim: .3f}'
else:
return ''
res[res < 0] = 0
res *= 255
res = res.astype(np.uint8)
res = res.transpose((1, 2, 0))
label = label.numpy()[0]
label *= 255
label = label.astype(np.uint8)
label = label.transpose((1, 2, 0))
if YCRCB_jpeg:
res = cv2.cvtColor(res, cv2.COLOR_RGB2YCR_CB)[:, :, 0]
label = cv2.cvtColor(label, cv2.COLOR_RGB2YCR_CB)[:, :, 0]
ave_psnr += psnr(res, label, data_range=255)
ave_ssim += ski_ssim(res,
label,
data_range=255,
multichannel=False)
ct_num += 1
else:
ave_psnr += psnr(res, label, data_range=255)
ave_ssim += ski_ssim(res, label, data_range=255, multichannel=True)
ct_num += 1
print('psnr_jpeg: ' + str(ave_psnr / float(ct_num)) + '.')
print('ssim_jpeg: ' + str(ave_ssim / float(ct_num)) + '.')
ave_psnr = 0.0
ave_ssim = 0.0
ct_num = 0
for test_iter, te_info in enumerate(
open(test3_list_pth).read().splitlines()):
te_pair_pth = testroot3 + te_info
for vi in range(v):
ct_num += 1
hazy_vi = hazy[bi, vi]
hazy_vi = hazy_vi.unsqueeze(dim=0)
hazy_vi = Variable(hazy_vi, requires_grad=False).cuda()
res = cleaner(hazy_vi)
res = res.data.cpu().numpy()[0]
res[res > 1] = 1
res[res < 0] = 0
res *= 255
res = res.astype(np.uint8)
res = res.transpose((1, 2, 0))
ave_psnr += psnr(res, label_v, data_range=255)
ave_ssim += ski_ssim(res,
label_v,
data_range=255,
multichannel=True)
Image.fromarray(res).save(show_dst + im_name[0].split('.')[0] +
'_' + str(vi + 1) + '.png')
elif data_name == 'DCPDNData':
ct_num += 1
label = label.numpy()[0]
label = label.transpose((1, 2, 0))
hazy = Variable(hazy, requires_grad=False).cuda()
res = cleaner(hazy)
res = res.data.cpu().numpy()[0]
res[res > 1] = 1
res[res < 0] = 0
res = res.transpose((1, 2, 0))
ave_psnr += psnr(res, label, data_range=1)
def __call__(self,
cgp,
gpuID,
epoch_num=150,
gpu_num=1,
result_file='results',
rank=4,
layer_num=1,
steps=4,
model_id=1):
print('GPUID :', gpuID)
print('epoch_num:', epoch_num)
# define model
torch.manual_seed(2018)
torch.cuda.manual_seed(2018)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True
L1_loss = nn.L1Loss()
L1_loss = L1_loss.cuda(gpuID)
if model_id == 1:
from model.OWAN import Network as Network1
model = Network1(16, layer_num, L1_loss, gpuID=gpuID, steps=steps)
elif model_id == 2:
from model.model_2order_ming import Network as Network2
model = Network2(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 3:
from model.model_3order_ming import Network as Network3
model = Network3(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 4:
from model.model_4order_ming import Network as Network4
model = Network4(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 5:
from model.model_2order_ming_unshare import Network as Network5
model = Network5(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 6:
from model.model_3order_ming_unshare import Network as Network6
model = Network6(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 7:
from model.model_4order_ming_unshare import Network as Network7
model = Network7(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 8:
from model.model_2order_tt import Network as Network8
model = Network8(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 9:
from model.model_3order_tt import Network as Network9
model = Network9(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 10:
from model.model_4order_tt import Network as Network10
model = Network10(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 11:
from model.model_2order_tr_share import Network as Network11
model = Network11(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 12:
from model.model_3order_tr_share import Network as Network12
model = Network12(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 13:
from model.model_4order_tr_share import Network as Network13
model = Network13(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 14:
from model.model_2order_tr_unshare import Network as Network14
model = Network14(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
elif model_id == 15:
from model.model_3order_tr_unshare import Network as Network15
model = Network15(16,
layer_num,
L1_loss,
gpuID=gpuID,
steps=steps,
rank=rank)
if gpu_num > 1:
device_ids = [i for i in range(gpu_num)]
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda(gpuID)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
print('Param:', utils.count_parameters_in_MB(model))
optimizer = optim.Adam(model.parameters(),
lr=0.001,
betas=(0.9, 0.999))
# for output images
if not os.path.exists(result_file):
os.makedirs('./' + result_file + '/Inputs')
os.makedirs('./' + result_file + '/Outputs')
os.makedirs('./' + result_file + '/Targets')
os.makedirs(result_file + '/Models')
checkpoint_path = None
model_dir = './' + result_file + '/Models/'
model_list = os.listdir(model_dir)
flag = 0
for num in range(100, 0, -1):
for model_name in model_list:
if str(num) in model_name:
checkpoint_path = model_dir + model_name
flag = 1
break
if flag == 1:
break
model, optimizer, last_epoch = load_checkpoint(model, checkpoint_path,
optimizer)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=epoch_num)
test_interval = 5
# Train loop
for epoch in range(1, epoch_num + 1):
scheduler.step()
if epoch <= last_epoch:
continue
start_time = time.time()
print('epoch', epoch)
train_loss = 0
for module in model.children():
module.train(True)
for ite, (input, target) in enumerate(self.dataloader):
lr_patch = Variable(input, requires_grad=False).cuda(gpuID)
hr_patch = Variable(target, requires_grad=False).cuda(gpuID)
optimizer.zero_grad()
output = model(lr_patch)
l1_loss = L1_loss(output, hr_patch)
l1_loss.backward()
optimizer.step()
train_loss += l1_loss.item()
if ite % 500 == 0:
vutils.save_image(lr_patch.data,
'./' + result_file +
'/input_sample%d.png' % gpuID,
normalize=False)
vutils.save_image(hr_patch.data,
'./' + result_file +
'/target_sample%d.png' % gpuID,
normalize=False)
vutils.save_image(output.data,
'./' + result_file +
'/output_sample%d.png' % gpuID,
normalize=False)
print('Train set : Average loss: {:.4f}'.format(train_loss))
print('time ', time.time() - start_time)
# check val/test performance
if epoch % test_interval == 0:
with torch.no_grad():
print('------------------------')
for module in model.children():
module.train(False)
# mild
test_psnr = 0
test_ssim = 0
eps = 1e-10
test_ite = 0
for i, (input, target) in enumerate(self.val_loader):
lr_patch = Variable(input.float(),
requires_grad=False).cuda(gpuID)
hr_patch = Variable(target.float(),
requires_grad=False).cuda(gpuID)
output = model(lr_patch)
# save images
vutils.save_image(output.data,
'./' + result_file +
'/Outputs/%05d.png' % (int(i)),
padding=0,
normalize=False)
vutils.save_image(lr_patch.data,
'./' + result_file +
'/Inputs/%05d.png' % (int(i)),
padding=0,
normalize=False)
vutils.save_image(hr_patch.data,
'./' + result_file +
'/Targets/%05d.png' % (int(i)),
padding=0,
normalize=False)
# Calculation of SSIM and PSNR values
output = output.data.cpu().numpy()[0]
output[output > 1] = 1
output[output < 0] = 0
output = output.transpose((1, 2, 0))
hr_patch = hr_patch.data.cpu().numpy()[0]
hr_patch[hr_patch > 1] = 1
hr_patch[hr_patch < 0] = 0
hr_patch = hr_patch.transpose((1, 2, 0))
# SSIM
test_ssim += ski_ssim(output,
hr_patch,
data_range=1,
multichannel=True)
# PSNR
imdf = (output - hr_patch)**2
mse = np.mean(imdf) + eps
test_psnr += 10 * math.log10(1.0 / mse)
test_ite += 1
test_psnr /= (test_ite)
test_ssim /= (test_ite)
print('Valid PSNR: {:.4f}'.format(test_psnr))
print('Valid SSIM: {:.4f}'.format(test_ssim))
print('------------------------')
torch.save(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, './' + result_file +
'/Models/model_%d.pth.tar' % int(epoch))
return train_loss
clean_F = cleaner(img_C, img_M, img_F, task)
img_C, img_M, img_F = ProcessResult(clean_F)
clean_F = clean_F.data.cpu().numpy()[0]
clean_F[clean_F > 1] = 1
clean_F[clean_F < 0] = 0
clean_F = clean_F * 255
clean_F = clean_F.astype(np.uint8)
clean_F = clean_F.transpose((1, 2, 0))
if YCRCB_rain:
clean_F = cv2.cvtColor(clean_F, cv2.COLOR_RGB2YCR_CB)[:, :, 0]
label_F = cv2.cvtColor(label_F, cv2.COLOR_RGB2YCR_CB)[:, :, 0]
ave_psnr += psnr(clean_F, label_F, data_range=255)
ave_ssim += ski_ssim(clean_F,
label_F,
data_range=255,
multichannel=False)
ct_num += 1.0
else:
ave_psnr += psnr(clean_F, label_F, data_range=255)
ave_ssim += ski_ssim(clean_F,
label_F,
data_range=255,
multichannel=True)
ct_num += 1.0
print('Rain-PSNR: ' + str(ave_psnr / ct_num) + '.')
print('Rain-SSIM: ' + str(ave_ssim / ct_num) + '.')
# Jpeg
ave_psnr = 0.0
ave_ssim = 0.0
def __call__(self, cgp, gpuID, epoch_num=150, gpu_num=1):
print('GPUID :', gpuID)
print('epoch_num:', epoch_num)
# define model
torch.manual_seed(2018)
torch.cuda.manual_seed(2018)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True
L1_loss = nn.L1Loss()
L1_loss = L1_loss.cuda(gpuID)
model = Network(16, 10, L1_loss, gpuID=gpuID)
if gpu_num > 1:
device_ids = [i for i in range(gpu_num)]
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda(gpuID)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
print('Param:', utils.count_parameters_in_MB(model))
optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epoch_num)
test_interval = 5
# for output images
if not os.path.exists('./results'):
os.makedirs('./results/Inputs')
os.makedirs('./results/Outputs')
os.makedirs('./results/Targets')
# Train loop
for epoch in range(1, epoch_num+1):
scheduler.step()
start_time = time.time()
print('epoch', epoch)
train_loss = 0
for module in model.children():
module.train(True)
for ite, (input, target) in enumerate(self.dataloader):
lr_patch = Variable(input, requires_grad=False).cuda(gpuID)
hr_patch = Variable(target, requires_grad=False).cuda(gpuID)
optimizer.zero_grad()
output = model(lr_patch)
l1_loss = L1_loss(output, hr_patch)
l1_loss.backward()
optimizer.step()
train_loss += l1_loss.item()
if ite % 500 == 0:
vutils.save_image(lr_patch.data, './input_sample%d.png' % gpuID, normalize=False)
vutils.save_image(hr_patch.data, './target_sample%d.png' % gpuID, normalize=False)
vutils.save_image(output.data, './output_sample%d.png' % gpuID, normalize=False)
print('Train set : Average loss: {:.4f}'.format(train_loss))
print('time ', time.time()-start_time)
# check val/test performance
if epoch % test_interval == 0:
with torch.no_grad():
print('------------------------')
for module in model.children():
module.train(False)
test_psnr = 0
test_ssim = 0
eps = 1e-10
test_ite = 0
for _, (input, target) in enumerate(self.val_loader):
lr_patch = Variable(input, requires_grad=False).cuda(gpuID)
hr_patch = Variable(target, requires_grad=False).cuda(gpuID)
output = model(lr_patch)
# save images
vutils.save_image(output.data, './results/Outputs/%05d.png' % (int(i)), padding=0, normalize=False)
vutils.save_image(lr_patch.data, './results/Inputs/%05d.png' % (int(i)), padding=0, normalize=False)
vutils.save_image(hr_patch.data, './results/Targets/%05d.png' % (int(i)), padding=0, normalize=False)
# Calculation of SSIM and PSNR values
output = output.data.cpu().numpy()[0]
output[output>1] = 1
output[output<0] = 0
output = output.transpose((1,2,0))
hr_patch = hr_patch.data.cpu().numpy()[0]
hr_patch[hr_patch>1] = 1
hr_patch[hr_patch<0] = 0
hr_patch = hr_patch.transpose((1,2,0))
# SSIM
test_ssim+= ski_ssim(output, hr_patch, data_range=1, multichannel=True)
# PSNR
imdf = (output - hr_patch) ** 2
mse = np.mean(imdf) + eps
test_psnr+= 10 * math.log10(1.0/mse)
test_ite += 1
test_psnr /= (test_ite)
test_ssim /= (test_ite)
print('Valid PSNR: {:.4f}'.format(test_psnr))
print('Valid SSIM: {:.4f}'.format(test_ssim))
f = open('PSNR.txt', 'a')
writer = csv.writer(f, lineterminator='\n')
writer.writerow([epoch, test_psnr, test_ssim])
f.close()
print('------------------------')
torch.save(model.state_dict(), './model_%d.pth' % int(epoch))
return train_loss
def ssim(image, reference_image):
r"""Compute the score of ``image`` regarding ``reference_image``
with the *Structural Similarity Index Measure* (SSIM) metric.
See [1]_, [2]_, [3]_ and [4]_ for more information.
The SSIM index is calculated on various windows of an image.
The measure between two windows :math:`x` and :math:`y` of common size
:math:`N.N` is:
.. math::
\hbox{SSIM}(x,y) = \frac{(2\mu_x\mu_y + c_1)(2\sigma_{xy} + c_2)}{(\mu_x^2 + \mu_y^2 + c_1)(\sigma_x^2 + \sigma_y^2 + c_2)}
with:
* :math:`\scriptstyle\mu_x` the average of :math:`\scriptstyle x`;
* :math:`\scriptstyle\mu_y` the average of :math:`\scriptstyle y`;
* :math:`\scriptstyle\sigma_x^2` the variance of :math:`\scriptstyle x`;
* :math:`\scriptstyle\sigma_y^2` the variance of :math:`\scriptstyle y`;
* :math:`\scriptstyle \sigma_{xy}` the covariance of :math:`\scriptstyle x` and :math:`\scriptstyle y`;
* :math:`\scriptstyle c_1 = (k_1L)^2`, :math:`\scriptstyle c_2 = (k_2L)^2` two variables to stabilize the division with weak denominator;
* :math:`\scriptstyle L` the dynamic range of the pixel-values (typically this is :math:`\scriptstyle 2^{\#bits\ per\ pixel}-1`);
* :math:`\scriptstyle k_1 = 0.01` and :math:`\scriptstyle k_2 = 0.03` by default.
The SSIM index satisfies the condition of symmetry:
.. math::
\text{SSIM}(x, y) = \text{SSIM}(y, x)
Parameters
----------
image: 2D ndarray
The cleaned image returned by the image cleanning algorithm to assess.
reference_image: 2D ndarray
The actual clean image (the best result that can be expected for the
image cleaning algorithm).
Returns
-------
float
The score of the image cleaning algorithm for the given image.
References
----------
.. [1] Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P.
(2004). Image quality assessment: From error visibility to
structural similarity. IEEE Transactions on Image Processing,
13, 600-612.
https://ece.uwaterloo.ca/~z70wang/publications/ssim.pdf,
DOI:10.1.1.11.2477
.. [2] Avanaki, A. N. (2009). Exact global histogram specification
optimized for structural similarity. Optical Review, 16, 613-621.
http://arxiv.org/abs/0901.0065,
DOI:10.1007/s10043-009-0119-z
.. [3] http://scikit-image.org/docs/dev/api/skimage.measure.html#compare-ssim
.. [4] https://en.wikipedia.org/wiki/Structural_similarity
"""
# Copy and cast images to prevent tricky bugs
# See https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.astype.html#numpy-ndarray-astype
image = image.astype('float64', copy=True)
reference_image = reference_image.astype('float64', copy=True)
# TODO: the following two lines may be wrong...
image[np.isnan(image)] = 0
reference_image[np.isnan(reference_image)] = 0
ssim_val, ssim_image = ski_ssim(image,
reference_image,
full=True,
gaussian_weights=True,
sigma=0.5)
return float(ssim_val)
for i, (input, target) in enumerate(test_dataloader):
lr_patch = Variable(input.float(), requires_grad=False).cuda(gpuID)
hr_patch = Variable(target.float(), requires_grad=False).cuda(gpuID)
output = model(lr_patch)
# Calculation of SSIM and PSNR values
output = output.data.cpu().numpy()
output[output > 1] = 1
output[output < 0] = 0
output = output.transpose((0, 2, 3, 1))
hr_patch = hr_patch.data.cpu().numpy()
hr_patch[hr_patch > 1] = 1
hr_patch[hr_patch < 0] = 0
hr_patch = hr_patch.transpose((0, 2, 3, 1))
# SSIM
for index in range(output.shape[0]):
test_ssim += ski_ssim(output[index],
hr_patch[index],
data_range=1,
multichannel=True)
# PSNR
for index in range(output.shape[0]):
imdf = (output[index] - hr_patch[index])**2
mse = np.mean(imdf) + eps
test_psnr += 10 * math.log10(1.0 / mse)
test_ite += 1
test_psnr /= (test_ite)
test_ssim /= (test_ite)
print('Test PSNR: {:.4f}'.format(test_psnr))
print('Test SSIM: {:.4f}'.format(test_ssim))
print('------------------------')
normalize=False)
vutils.save_image(lr_patch.data,
'./results/Inputs/%05d.png' % (int(i)),
padding=0,
normalize=False)
vutils.save_image(hr_patch.data,
'./results/Targets/%05d.png' % (int(i)),
padding=0,
normalize=False)
# SSIM and PSNR
output = output.data.cpu().numpy()[0]
output[output > 1] = 1
output[output < 0] = 0
output = output.transpose((1, 2, 0))
hr_patch = hr_patch.data.cpu().numpy()[0]
hr_patch[hr_patch > 1] = 1
hr_patch[hr_patch < 0] = 0
hr_patch = hr_patch.transpose((1, 2, 0))
# SSIM
test_ssim += ski_ssim(output, hr_patch, data_range=1, multichannel=True)
# PSNR
imdf = (output - hr_patch)**2
mse = np.mean(imdf) + eps
test_psnr += 10 * math.log10(1.0 / mse)
test_ite += 1
test_psnr /= (test_ite)
test_ssim /= (test_ite)
print('Test PSNR: {:.4f}'.format(test_psnr))
print('Test SSIM: {:.4f}'.format(test_ssim))
print('------------------------')
