415 lines
12 KiB
Python
Executable File
415 lines
12 KiB
Python
Executable File
import os
|
|
import shutil
|
|
import sys
|
|
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.utils as vutils
|
|
from skimage import color, io
|
|
from torch.autograd import Variable
|
|
|
|
cv2.setNumThreads(0)
|
|
|
|
# l: [-50,50]
|
|
# ab: [-128, 128]
|
|
l_norm, ab_norm = 1.0, 1.0
|
|
l_mean, ab_mean = 50.0, 0
|
|
|
|
|
|
###### utility ######
|
|
def to_np(x):
|
|
return x.data.cpu().numpy()
|
|
|
|
|
|
def utf8_str(in_str):
|
|
try:
|
|
in_str = in_str.decode("UTF-8")
|
|
except Exception:
|
|
in_str = in_str.encode("UTF-8").decode("UTF-8")
|
|
return in_str
|
|
|
|
|
|
class MovingAvg(object):
|
|
def __init__(self, pool_size=100):
|
|
from queue import Queue
|
|
|
|
self.pool = Queue(maxsize=pool_size)
|
|
self.sum = 0
|
|
self.curr_pool_size = 0
|
|
|
|
def set_curr_val(self, val):
|
|
if not self.pool.full():
|
|
self.curr_pool_size += 1
|
|
self.pool.put_nowait(val)
|
|
else:
|
|
last_first_val = self.pool.get_nowait()
|
|
self.pool.put_nowait(val)
|
|
self.sum -= last_first_val
|
|
|
|
self.sum += val
|
|
return self.sum / self.curr_pool_size
|
|
|
|
|
|
###### image normalization ######
|
|
def center_l(l):
|
|
# normalization for l
|
|
l_mc = (l - l_mean) / l_norm
|
|
return l_mc
|
|
|
|
|
|
# denormalization for l
|
|
def uncenter_l(l):
|
|
return l * l_norm + l_mean
|
|
|
|
|
|
# normalization for ab
|
|
def center_ab(ab):
|
|
return (ab - ab_mean) / ab_norm
|
|
|
|
|
|
# normalization for lab image
|
|
def center_lab_img(img_lab):
|
|
return (
|
|
img_lab / np.array((l_norm, ab_norm, ab_norm))[:, np.newaxis, np.newaxis]
|
|
- np.array((l_mean / l_norm, ab_mean / ab_norm, ab_mean / ab_norm))[:, np.newaxis, np.newaxis]
|
|
)
|
|
|
|
|
|
###### color space transformation ######
|
|
def rgb2lab_transpose(img_rgb):
|
|
return color.rgb2lab(img_rgb).transpose((2, 0, 1))
|
|
|
|
|
|
def lab2rgb(img_l, img_ab):
|
|
"""INPUTS
|
|
img_l XxXx1 [0,100]
|
|
img_ab XxXx2 [-100,100]
|
|
OUTPUTS
|
|
returned value is XxXx3"""
|
|
pred_lab = np.concatenate((img_l, img_ab), axis=2).astype("float64")
|
|
pred_rgb = color.lab2rgb(pred_lab)
|
|
pred_rgb = (np.clip(pred_rgb, 0, 1) * 255).astype("uint8")
|
|
return pred_rgb
|
|
|
|
|
|
def gray2rgb_batch(l):
|
|
# gray image tensor to rgb image tensor
|
|
l_uncenter = uncenter_l(l)
|
|
l_uncenter = l_uncenter / (2 * l_mean)
|
|
return torch.cat((l_uncenter, l_uncenter, l_uncenter), dim=1)
|
|
|
|
|
|
def lab2rgb_transpose(img_l, img_ab):
|
|
"""INPUTS
|
|
img_l 1xXxX [0,100]
|
|
img_ab 2xXxX [-100,100]
|
|
OUTPUTS
|
|
returned value is XxXx3"""
|
|
pred_lab = np.concatenate((img_l, img_ab), axis=0).transpose((1, 2, 0))
|
|
return (np.clip(color.lab2rgb(pred_lab), 0, 1) * 255).astype("uint8")
|
|
|
|
|
|
def lab2rgb_transpose_mc(img_l_mc, img_ab_mc):
|
|
if isinstance(img_l_mc, Variable):
|
|
img_l_mc = img_l_mc.data.cpu()
|
|
if isinstance(img_ab_mc, Variable):
|
|
img_ab_mc = img_ab_mc.data.cpu()
|
|
|
|
if img_l_mc.is_cuda:
|
|
img_l_mc = img_l_mc.cpu()
|
|
if img_ab_mc.is_cuda:
|
|
img_ab_mc = img_ab_mc.cpu()
|
|
|
|
assert img_l_mc.dim() == 3 and img_ab_mc.dim() == 3, "only for batch input"
|
|
|
|
img_l = img_l_mc * l_norm + l_mean
|
|
img_ab = img_ab_mc * ab_norm + ab_mean
|
|
pred_lab = torch.cat((img_l, img_ab), dim=0)
|
|
grid_lab = pred_lab.numpy().astype("float64")
|
|
return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
|
|
|
|
|
def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8):
|
|
if isinstance(img_l_mc, Variable):
|
|
img_l_mc = img_l_mc.data.cpu()
|
|
if isinstance(img_ab_mc, Variable):
|
|
img_ab_mc = img_ab_mc.data.cpu()
|
|
|
|
if img_l_mc.is_cuda:
|
|
img_l_mc = img_l_mc.cpu()
|
|
if img_ab_mc.is_cuda:
|
|
img_ab_mc = img_ab_mc.cpu()
|
|
|
|
assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input"
|
|
|
|
img_l = img_l_mc * l_norm + l_mean
|
|
img_ab = img_ab_mc * ab_norm + ab_mean
|
|
pred_lab = torch.cat((img_l, img_ab), dim=1)
|
|
grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64")
|
|
return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
|
|
|
|
|
###### loss functions ######
|
|
def feature_normalize(feature_in):
|
|
feature_in_norm = torch.norm(feature_in, 2, 1, keepdim=True) + sys.float_info.epsilon
|
|
feature_in_norm = torch.div(feature_in, feature_in_norm)
|
|
return feature_in_norm
|
|
|
|
|
|
def statistics_matching(feature1, feature2):
|
|
N, C, H, W = feature1.shape
|
|
feature1 = feature1.view(N, C, -1)
|
|
feature2 = feature2.view(N, C, -1)
|
|
|
|
mean1 = feature1.mean(dim=-1)
|
|
mean2 = feature2.mean(dim=-1)
|
|
std1 = feature1.var(dim=-1).sqrt()
|
|
std2 = feature2.var(dim=-1).sqrt()
|
|
|
|
return mse_loss(mean1, mean2) + mse_loss(std1, std2)
|
|
|
|
|
|
def cosine_similarity(input, target):
|
|
input_norm = torch.norm(input, 2, 1, keepdim=True) + sys.float_info.epsilon
|
|
target_norm = torch.norm(target, 2, 1, keepdim=True) + sys.float_info.epsilon
|
|
normalized_input = torch.div(input, input_norm)
|
|
normalized_target = torch.div(target, target_norm)
|
|
cos_similarity = torch.mul(normalized_input, normalized_target)
|
|
return torch.sum(cos_similarity, dim=1, keepdim=True)
|
|
|
|
|
|
def mse_loss(input, target=0):
|
|
return torch.mean((input - target) ** 2)
|
|
|
|
|
|
def l1_loss(input, target=0):
|
|
return torch.mean(torch.abs(input - target))
|
|
|
|
|
|
def calc_ab_gradient(input_ab):
|
|
x_grad = input_ab[:, :, :, 1:] - input_ab[:, :, :, :-1]
|
|
y_grad = input_ab[:, :, 1:, :] - input_ab[:, :, :-1, :]
|
|
return x_grad, y_grad
|
|
|
|
|
|
def calc_tv_loss(input):
|
|
x_grad = input[:, :, :, 1:] - input[:, :, :, :-1]
|
|
y_grad = input[:, :, 1:, :] - input[:, :, :-1, :]
|
|
return torch.sum(x_grad ** 2) / x_grad.nelement() + torch.sum(y_grad ** 2) / y_grad.nelement()
|
|
|
|
|
|
def calc_cosine_dist_loss(input, target):
|
|
input_norm = torch.norm(input, 2, 1, keepdim=True) + sys.float_info.epsilon
|
|
target_norm = torch.norm(target, 2, 1, keepdim=True) + sys.float_info.epsilon
|
|
normalized_input = torch.div(input, input_norm)
|
|
normalized_target = torch.div(target, target_norm)
|
|
cos_dist = torch.mul(normalized_input, normalized_target)
|
|
return torch.mean(1 - torch.sum(cos_dist, dim=1))
|
|
|
|
|
|
def weighted_mse_loss(input, target, weights):
|
|
out = (input - target) ** 2
|
|
out = out * weights.expand_as(out)
|
|
return out.mean()
|
|
|
|
|
|
def weighted_l1_loss(input, target, weights):
|
|
out = torch.abs(input - target)
|
|
out = out * weights.expand_as(out)
|
|
return out.mean()
|
|
|
|
|
|
def colorfulness(input_ab):
|
|
"""
|
|
according to the paper: Measuring colourfulness in natural images
|
|
input is batches of ab tensors in lab space
|
|
"""
|
|
N, C, H, W = input_ab.shape
|
|
a = input_ab[:, 0:1, :, :]
|
|
b = input_ab[:, 1:2, :, :]
|
|
|
|
a = a.view(N, -1)
|
|
b = b.view(N, -1)
|
|
|
|
sigma_a = torch.std(a, dim=-1)
|
|
sigma_b = torch.std(b, dim=-1)
|
|
|
|
mean_a = torch.mean(a, dim=-1)
|
|
mean_b = torch.mean(b, dim=-1)
|
|
|
|
return torch.sqrt(sigma_a ** 2 + sigma_b ** 2) + 0.37 * torch.sqrt(mean_a ** 2 + mean_b ** 2)
|
|
|
|
|
|
###### video related #######
|
|
def save_frames(image, image_folder, index=None, image_name=None):
|
|
if image is not None:
|
|
image = np.clip(image, 0, 255).astype(np.uint8)
|
|
if image_name:
|
|
io.imsave(os.path.join(image_folder, image_name), image)
|
|
else:
|
|
io.imsave(os.path.join(image_folder, str(index).zfill(5) + ".jpg"), image)
|
|
|
|
|
|
def folder2vid(image_folder, output_dir, filename):
|
|
images = [img for img in os.listdir(image_folder) if img.endswith(".jpg") or img.endswith(".png")]
|
|
images.sort()
|
|
frame = cv2.imread(os.path.join(image_folder, images[0]))
|
|
height, width, layers = frame.shape
|
|
print("writing to video " + os.path.join(output_dir, filename))
|
|
video = cv2.VideoWriter(
|
|
os.path.join(output_dir, filename), cv2.VideoWriter_fourcc("D", "I", "V", "X"), 24, (width, height)
|
|
)
|
|
|
|
for image in images:
|
|
video.write(cv2.imread(os.path.join(image_folder, image)))
|
|
|
|
video.release()
|
|
|
|
# import imageio
|
|
# frames = []
|
|
# for image in images:
|
|
# frames.append(imageio.imread(os.path.join(image_folder, image)))
|
|
# imageio.mimsave('movie.gif', frames)
|
|
|
|
|
|
###### file system ######
|
|
def get_size(start_path="."):
|
|
total_size = 0
|
|
for dirpath, dirnames, filenames in os.walk(start_path):
|
|
for f in filenames:
|
|
fp = os.path.join(dirpath, f)
|
|
total_size += os.path.getsize(fp)
|
|
return total_size
|
|
|
|
|
|
def mkdir_if_not(dir_path):
|
|
if not os.path.exists(dir_path):
|
|
os.makedirs(dir_path)
|
|
|
|
|
|
def parse(parser, save=True):
|
|
opt = parser.parse_args(args=[])
|
|
args = vars(opt)
|
|
|
|
from time import gmtime, strftime
|
|
|
|
print("------------ Options -------------")
|
|
for k, v in sorted(args.items()):
|
|
print("%s: %s" % (str(k), str(v)))
|
|
print("-------------- End ----------------")
|
|
|
|
# save to the disk
|
|
if save:
|
|
file_name = os.path.join("opt.txt")
|
|
with open(file_name, "wt") as opt_file:
|
|
opt_file.write(os.path.basename(sys.argv[0]) + " " + strftime("%Y-%m-%d %H:%M:%S", gmtime()) + "\n")
|
|
opt_file.write("------------ Options -------------\n")
|
|
for k, v in sorted(args.items()):
|
|
opt_file.write("%s: %s\n" % (str(k), str(v)))
|
|
opt_file.write("-------------- End ----------------\n")
|
|
return opt
|
|
|
|
|
|
###### interactive ######
|
|
def clean_tensorboard(directory):
|
|
folder_list = os.walk(directory).__next__()[1]
|
|
for folder in folder_list:
|
|
folder = directory + folder
|
|
if get_size(folder) < 10000000:
|
|
print("delete the folder of " + folder)
|
|
shutil.rmtree(folder)
|
|
|
|
|
|
def imshow(input_image, title=None, type_conversion=False):
|
|
inp = input_image
|
|
if type_conversion or type(input_image) is torch.Tensor:
|
|
inp = input_image.numpy()
|
|
else:
|
|
inp = input_image
|
|
fig = plt.figure()
|
|
if inp.ndim == 2:
|
|
fig = plt.imshow(inp, cmap="gray", clim=[0, 255])
|
|
else:
|
|
fig = plt.imshow(np.transpose(inp, [1, 2, 0]).astype(np.uint8))
|
|
plt.axis("off")
|
|
fig.axes.get_xaxis().set_visible(False)
|
|
fig.axes.get_yaxis().set_visible(False)
|
|
plt.title(title)
|
|
|
|
|
|
def imshow_lab(input_lab):
|
|
plt.imshow((batch_lab2rgb_transpose_mc(input_lab[:32, 0:1, :, :], input_lab[:32, 1:3, :, :])).astype(np.uint8))
|
|
|
|
|
|
###### vgg preprocessing ######
|
|
def vgg_preprocess(tensor):
|
|
# input is RGB tensor which ranges in [0,1]
|
|
# output is BGR tensor which ranges in [0,255]
|
|
tensor_bgr = torch.cat((tensor[:, 2:3, :, :], tensor[:, 1:2, :, :], tensor[:, 0:1, :, :]), dim=1)
|
|
tensor_bgr_ml = tensor_bgr - torch.Tensor([0.40760392, 0.45795686, 0.48501961]).type_as(tensor_bgr).view(1, 3, 1, 1)
|
|
return tensor_bgr_ml * 255
|
|
|
|
|
|
def torch_vgg_preprocess(tensor):
|
|
# pytorch version normalization
|
|
# note that both input and output are RGB tensors;
|
|
# input and output ranges in [0,1]
|
|
# normalize the tensor with mean and variance
|
|
tensor_mc = tensor - torch.Tensor([0.485, 0.456, 0.406]).type_as(tensor).view(1, 3, 1, 1)
|
|
return tensor_mc / torch.Tensor([0.229, 0.224, 0.225]).type_as(tensor_mc).view(1, 3, 1, 1)
|
|
|
|
|
|
def network_gradient(net, gradient_on=True):
|
|
for param in net.parameters():
|
|
param.requires_grad = bool(gradient_on)
|
|
return net
|
|
|
|
|
|
##### color space
|
|
xyz_from_rgb = np.array(
|
|
[[0.412453, 0.357580, 0.180423], [0.212671, 0.715160, 0.072169], [0.019334, 0.119193, 0.950227]]
|
|
)
|
|
rgb_from_xyz = np.array(
|
|
[[3.24048134, -0.96925495, 0.05564664], [-1.53715152, 1.87599, -0.20404134], [-0.49853633, 0.04155593, 1.05731107]]
|
|
)
|
|
|
|
|
|
def tensor_lab2rgb(input):
|
|
"""
|
|
n * 3* h *w
|
|
"""
|
|
input_trans = input.transpose(1, 2).transpose(2, 3) # n * h * w * 3
|
|
L, a, b = input_trans[:, :, :, 0:1], input_trans[:, :, :, 1:2], input_trans[:, :, :, 2:]
|
|
y = (L + 16.0) / 116.0
|
|
x = (a / 500.0) + y
|
|
z = y - (b / 200.0)
|
|
|
|
neg_mask = z.data < 0
|
|
z[neg_mask] = 0
|
|
xyz = torch.cat((x, y, z), dim=3)
|
|
|
|
mask = xyz.data > 0.2068966
|
|
mask_xyz = xyz.clone()
|
|
mask_xyz[mask] = torch.pow(xyz[mask], 3.0)
|
|
mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.0) / 7.787
|
|
mask_xyz[:, :, :, 0] = mask_xyz[:, :, :, 0] * 0.95047
|
|
mask_xyz[:, :, :, 2] = mask_xyz[:, :, :, 2] * 1.08883
|
|
|
|
rgb_trans = torch.mm(mask_xyz.view(-1, 3), torch.from_numpy(rgb_from_xyz).type_as(xyz)).view(
|
|
input.size(0), input.size(2), input.size(3), 3
|
|
)
|
|
rgb = rgb_trans.transpose(2, 3).transpose(1, 2)
|
|
|
|
mask = rgb > 0.0031308
|
|
mask_rgb = rgb.clone()
|
|
mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055
|
|
mask_rgb[~mask] = rgb[~mask] * 12.92
|
|
|
|
neg_mask = mask_rgb.data < 0
|
|
large_mask = mask_rgb.data > 1
|
|
mask_rgb[neg_mask] = 0
|
|
mask_rgb[large_mask] = 1
|
|
return mask_rgb
|