272 lines
8.6 KiB
Python
272 lines
8.6 KiB
Python
"""This module contains simple helper functions """
|
|
from __future__ import print_function
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import scipy.ndimage as ndimage
|
|
import os
|
|
|
|
import cv2
|
|
|
|
import torch.nn.functional as F
|
|
|
|
def synthesize_fog(J, t, A=None):
|
|
"""
|
|
Synthesize hazy image base on optical model
|
|
I = J * t + A * (1 - t)
|
|
"""
|
|
|
|
if A is None:
|
|
A = 1
|
|
|
|
return J * t + A * (1 - t)
|
|
|
|
def reverse_fog(I, t, A=1, t0=0.01):
|
|
"""
|
|
Recover haze-free image using hazy image and depth
|
|
J = (I - A) / max(t, t0) + A
|
|
"""
|
|
|
|
t_clamp = torch.clamp(t, t0, 1)
|
|
J = (I-A) / t_clamp + A
|
|
return torch.clamp(J, -1, 1)
|
|
|
|
|
|
def fuse_images(real_I, rec_J, refine_J):
|
|
"""
|
|
real_I, rec_J, and refine_J: Images with shape hxwx3
|
|
"""
|
|
# realness features
|
|
mat_RGB2YMN = np.array([[0.299,0.587,0.114],
|
|
[0.30,0.04,-0.35],
|
|
[0.34,-0.6,0.17]])
|
|
|
|
recH,recW,recChl = rec_J.shape
|
|
rec_J_flat = rec_J.reshape([recH*recW,recChl])
|
|
rec_J_flat_YMN = (mat_RGB2YMN.dot(rec_J_flat.T)).T
|
|
rec_J_YMN = rec_J_flat_YMN.reshape(rec_J.shape)
|
|
|
|
refine_J_flat = refine_J.reshape([recH*recW,recChl])
|
|
refine_J_flat_YMN = (mat_RGB2YMN.dot(refine_J_flat.T)).T
|
|
refine_J_YMN = refine_J_flat_YMN.reshape(refine_J.shape)
|
|
|
|
real_I_flat = real_I.reshape([recH*recW,recChl])
|
|
real_I_flat_YMN = (mat_RGB2YMN.dot(real_I_flat.T)).T
|
|
real_I_YMN = real_I_flat_YMN.reshape(real_I.shape)
|
|
|
|
# gradient features
|
|
rec_Gx = cv2.Sobel(rec_J_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
|
|
rec_Gy = cv2.Sobel(rec_J_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
|
|
rec_GM = np.sqrt(rec_Gx**2 + rec_Gy**2)
|
|
|
|
refine_Gx = cv2.Sobel(refine_J_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
|
|
refine_Gy = cv2.Sobel(refine_J_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
|
|
refine_GM = np.sqrt(refine_Gx**2 + refine_Gy**2)
|
|
|
|
real_Gx = cv2.Sobel(real_I_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
|
|
real_Gy = cv2.Sobel(real_I_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
|
|
real_GM = np.sqrt(real_Gx**2 + real_Gy**2)
|
|
|
|
# similarity
|
|
rec_S_V = (2*real_GM*rec_GM+160)/(real_GM**2+rec_GM**2+160)
|
|
rec_S_M = (2*rec_J_YMN[:,:,1]*real_I_YMN[:,:,1]+130)/(rec_J_YMN[:,:,1]**2+real_I_YMN[:,:,1]**2+130)
|
|
rec_S_N = (2*rec_J_YMN[:,:,2]*real_I_YMN[:,:,2]+130)/(rec_J_YMN[:,:,2]**2+real_I_YMN[:,:,2]**2+130)
|
|
rec_S_R = (rec_S_M*rec_S_N).reshape([recH,recW])
|
|
|
|
refine_S_V = (2*real_GM*refine_GM+160)/(real_GM**2+refine_GM**2+160)
|
|
refine_S_M = (2*refine_J_YMN[:,:,1]*real_I_YMN[:,:,1]+130)/(refine_J_YMN[:,:,1]**2+real_I_YMN[:,:,1]**2+130)
|
|
refine_S_N = (2*refine_J_YMN[:,:,2]*real_I_YMN[:,:,2]+130)/(refine_J_YMN[:,:,2]**2+real_I_YMN[:,:,2]**2+130)
|
|
refine_S_R = (refine_S_M*refine_S_N).reshape([recH,recW])
|
|
|
|
|
|
rec_S = rec_S_R*np.power(rec_S_V, 0.4)
|
|
refine_S = refine_S_R*np.power(refine_S_V, 0.4)
|
|
|
|
|
|
fuseWeight = np.exp(rec_S)/(np.exp(rec_S)+np.exp(refine_S))
|
|
fuseWeightMap = fuseWeight.reshape([recH,recW,1]).repeat(3,axis=2)
|
|
|
|
fuse_J = rec_J*fuseWeightMap + refine_J*(1-fuseWeightMap)
|
|
return fuse_J
|
|
|
|
|
|
|
|
def get_tensor_dark_channel(img, neighborhood_size):
|
|
shape = img.shape
|
|
if len(shape) == 4:
|
|
img_min = torch.min(img, dim=1)
|
|
img_dark = F.max_pool2d(img_min, kernel_size=neighborhood_size, stride=1)
|
|
else:
|
|
raise NotImplementedError('get_tensor_dark_channel is only for 4-d tensor [N*C*H*W]')
|
|
|
|
return img_dark
|
|
|
|
|
|
|
|
def array2Tensor(in_array, gpu_id=-1):
|
|
in_shape = in_array.shape
|
|
if len(in_shape) == 2:
|
|
in_array = in_array[:,:,np.newaxis]
|
|
|
|
arr_tmp = in_array.transpose([2,0,1])
|
|
arr_tmp = arr_tmp[np.newaxis,:]
|
|
|
|
if gpu_id >= 0:
|
|
return torch.tensor(arr_tmp.astype(float)).to(gpu_id)
|
|
else:
|
|
return torch.tensor(arr_tmp.astype(float))
|
|
|
|
|
|
def tensor2im(input_image, imtype=np.uint8):
|
|
""""Converts a Tensor array into a numpy image array.
|
|
|
|
Parameters:
|
|
input_image (tensor) -- the input image tensor array
|
|
imtype (type) -- the desired type of the converted numpy array
|
|
"""
|
|
if not isinstance(input_image, np.ndarray):
|
|
if isinstance(input_image, torch.Tensor): # get the data from a variable
|
|
image_tensor = input_image.data
|
|
else:
|
|
return input_image
|
|
image_numpy = image_tensor[0].cpu().float().numpy() # convert it into a numpy array
|
|
if image_numpy.shape[0] == 1: # grayscale to RGB
|
|
image_numpy = np.tile(image_numpy, (3, 1, 1))
|
|
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 # post-processing: tranpose and scaling
|
|
else: # if it is a numpy array, do nothing
|
|
image_numpy = input_image
|
|
return image_numpy.astype(imtype)
|
|
|
|
def rescale_tensor(input_tensor):
|
|
""""Converts a Tensor array into the Tensor array whose data are identical to the image's.
|
|
[height, width] not [width, height]
|
|
|
|
Parameters:
|
|
input_image (tensor) -- the input image tensor array
|
|
imtype (type) -- the desired type of the converted numpy array
|
|
"""
|
|
|
|
if isinstance(input_tensor, torch.Tensor):
|
|
input_tmp = input_tensor.cpu().float()
|
|
output_tmp = (input_tmp + 1) / 2.0 * 255.0
|
|
output_tmp = output_tmp.to(torch.uint8)
|
|
else:
|
|
return input_tensor
|
|
|
|
return output_tmp.to(torch.float32) / 255.0
|
|
|
|
# if not isinstance(input_image, np.ndarray):
|
|
# if isinstance(input_image, torch.Tensor): # get the data from a variable
|
|
# image_tensor = input_image.data
|
|
# else:
|
|
# return input_image
|
|
# image_numpy = image_tensor.cpu().float().numpy() # convert it into a numpy array
|
|
# image_numpy = (image_numpy + 1) / 2.0 * white_color # post-processing: tranpose and scaling
|
|
# else: # if it is a numpy array, do nothing
|
|
# image_numpy = input_image
|
|
# return torch.from_numpy(image_numpy)
|
|
|
|
def my_imresize(in_array, tar_size):
|
|
oh = in_array.shape[0]
|
|
ow = in_array.shape[1]
|
|
|
|
if len(tar_size) == 2:
|
|
h_ratio = tar_size[0]/oh
|
|
w_ratio = tar_size[1]/ow
|
|
elif len(tar_size) == 1:
|
|
h_ratio = tar_size
|
|
w_ratio = tar_size
|
|
|
|
if len(in_array.shape) == 3:
|
|
return ndimage.zoom(in_array, (h_ratio, w_ratio, 1), prefilter=False)
|
|
else:
|
|
return ndimage.zoom(in_array, (h_ratio, w_ratio), prefilter=False)
|
|
|
|
def psnr(img, ref, max_val=1):
|
|
if isinstance(img, torch.Tensor):
|
|
distImg = img.cpu().float().numpy()
|
|
elif isinstance(img, np.ndarray):
|
|
distImg = img.astype(float)
|
|
else:
|
|
distImg = np.array(img).astype(float)
|
|
|
|
if isinstance(ref, torch.Tensor):
|
|
refImg = ref.cpu().float().numpy()
|
|
elif isinstance(ref, np.ndarray):
|
|
refImg = ref.astype(float)
|
|
else:
|
|
refImg = np.array(ref).astype(float)
|
|
|
|
rmse = np.sqrt( ((distImg-refImg)**2).mean() )
|
|
# rmse = np.std(distImg-refImg) # keep the same with RESIDE's criterion
|
|
return 20*np.log10(max_val/rmse)
|
|
|
|
|
|
def diagnose_network(net, name='network'):
|
|
"""Calculate and print the mean of average absolute(gradients)
|
|
|
|
Parameters:
|
|
net (torch network) -- Torch network
|
|
name (str) -- the name of the network
|
|
"""
|
|
mean = 0.0
|
|
count = 0
|
|
for param in net.parameters():
|
|
if param.grad is not None:
|
|
mean += torch.mean(torch.abs(param.grad.data))
|
|
count += 1
|
|
if count > 0:
|
|
mean = mean / count
|
|
print(name)
|
|
print(mean)
|
|
|
|
|
|
def save_image(image_numpy, image_path):
|
|
"""Save a numpy image to the disk
|
|
|
|
Parameters:
|
|
image_numpy (numpy array) -- input numpy array
|
|
image_path (str) -- the path of the image
|
|
"""
|
|
image_pil = Image.fromarray(image_numpy)
|
|
image_pil.save(image_path)
|
|
|
|
|
|
def print_numpy(x, val=True, shp=False):
|
|
"""Print the mean, min, max, median, std, and size of a numpy array
|
|
|
|
Parameters:
|
|
val (bool) -- if print the values of the numpy array
|
|
shp (bool) -- if print the shape of the numpy array
|
|
"""
|
|
x = x.astype(np.float64)
|
|
if shp:
|
|
print('shape,', x.shape)
|
|
if val:
|
|
x = x.flatten()
|
|
print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
|
|
np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
|
|
|
|
|
|
def mkdirs(paths):
|
|
"""create empty directories if they don't exist
|
|
|
|
Parameters:
|
|
paths (str list) -- a list of directory paths
|
|
"""
|
|
if isinstance(paths, list) and not isinstance(paths, str):
|
|
for path in paths:
|
|
mkdir(path)
|
|
else:
|
|
mkdir(paths)
|
|
|
|
|
|
def mkdir(path):
|
|
"""create a single empty directory if it didn't exist
|
|
|
|
Parameters:
|
|
path (str) -- a single directory path
|
|
"""
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|