423 lines
13 KiB
Python
Executable File
423 lines
13 KiB
Python
Executable File
import numbers
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import time
|
|
|
|
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 ######
|
|
# normalization for l
|
|
def center_l(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):
|
|
ab_mc = (ab - ab_mean) / ab_norm
|
|
return ab_mc
|
|
|
|
|
|
# normalization for lab image
|
|
def center_lab_img(img_lab):
|
|
img_lab_mc = (
|
|
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]
|
|
)
|
|
return img_lab_mc
|
|
|
|
|
|
###### color space transformation ######
|
|
def rgb2lab_transpose(img_rgb):
|
|
img_lab = color.rgb2lab(img_rgb).transpose((2, 0, 1))
|
|
return img_lab
|
|
|
|
|
|
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))
|
|
pred_rgb = (np.clip(color.lab2rgb(pred_lab), 0, 1) * 255).astype("uint8")
|
|
return pred_rgb
|
|
|
|
|
|
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")
|
|
grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
|
return grid_rgb
|
|
|
|
|
|
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")
|
|
grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
|
return grid_rgb
|
|
|
|
|
|
###### 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, :]
|
|
loss = torch.sum(x_grad ** 2) / x_grad.nelement() + torch.sum(y_grad ** 2) / y_grad.nelement()
|
|
return loss
|
|
|
|
|
|
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)
|
|
loss = out.mean()
|
|
return loss
|
|
|
|
|
|
def weighted_l1_loss(input, target, weights):
|
|
out = torch.abs(input - target)
|
|
out = out * weights.expand_as(out)
|
|
loss = out.mean()
|
|
return loss
|
|
|
|
|
|
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)
|
|
|
|
colorfulness = torch.sqrt(sigma_a ** 2 + sigma_b ** 2) + 0.37 * torch.sqrt(mean_a ** 2 + mean_b ** 2)
|
|
return colorfulness
|
|
|
|
|
|
###### 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")]
|
|
images.sort()
|
|
print(images)
|
|
frame = cv2.imread(os.path.join(image_folder, images[0]))
|
|
height, width, layers = frame.shape
|
|
# video = cv2.VideoWriter(os.path.join(output_dir, filename), 0x00000021, 24, (width, height))
|
|
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)
|
|
)
|
|
# video = cv2.VideoWriter(os.path.join(output_dir, filename), cv2.VideoWriter_fourcc('H', '2', '6', '4'), 24, (width, height))
|
|
for image in images:
|
|
video.write(cv2.imread(os.path.join(image_folder, image)))
|
|
|
|
# cv2.destroyAllWindows()
|
|
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 = tensor[:, [2, 1, 0], ...]
|
|
tensor_bgr_ml = tensor_bgr - torch.Tensor([0.40760392, 0.45795686, 0.48501961]).type_as(tensor_bgr).view(1, 3, 1, 1)
|
|
tensor_rst = tensor_bgr_ml * 255
|
|
return tensor_rst
|
|
|
|
|
|
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)
|
|
tensor_mc_norm = tensor_mc / torch.Tensor([0.229, 0.224, 0.225]).type_as(tensor_mc).view(1, 3, 1, 1)
|
|
return tensor_mc_norm
|
|
|
|
|
|
def network_gradient(net, gradient_on=True):
|
|
if gradient_on:
|
|
for param in net.parameters():
|
|
param.requires_grad = True
|
|
else:
|
|
for param in net.parameters():
|
|
param.requires_grad = False
|
|
return net
|
|
|
|
|
|
# def load_gray_image(img_path):
|
|
# img_path = utf8_str(img_path)
|
|
# img_gray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
|
|
# return img_gray
|
|
|
|
# def load_rgb_image(img_path):
|
|
# img_path = utf8_str(img_path)
|
|
# im = cv2.imread(img_path, cv2.IMREAD_COLOR)
|
|
# if im.shape[2] == 1:
|
|
# print('find gray images: ', img_path)
|
|
# im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
|
|
# img_rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
|
# return img_rgb
|
|
|
|
# def resize_img_longside(img_rgb, long_size, interpolation):
|
|
# h, w = img_rgb.shape[:2]
|
|
# scalar = long_size / max(h, w)
|
|
# resized_img_rgb = cv2.resize(img_rgb, (int(w * scalar), int(h * scalar)), interpolation=interpolation)
|
|
# return resized_img_rgb
|
|
|
|
# def resize_img_shortside(img_rgb, short_size, interpolation):
|
|
# h, w = img_rgb.shape[:2]
|
|
# scalar = short_size / min(h, w)
|
|
# resized_img_rgb = cv2.resize(img_rgb, (int(w * scalar), int(h * scalar)), interpolation=interpolation)
|
|
# return resized_img_rgb
|
|
|
|
# def resize_img(img_rgb, dsize, interpolation):
|
|
# if isinstance(dsize, numbers.Number):
|
|
# dsize = (int(dsize), int(dsize))
|
|
# resized_img_rgb = cv2.resize(img_rgb, dsize, interpolation=interpolation)
|
|
# return resized_img_rgb
|