init
This commit is contained in:
zhangmozhe
Executable
@@ -0,0 +1,154 @@
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import lib.functional as F
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from scipy.misc import imread, imresize, imsave
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from skimage import color
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
|
||||
# %% Dataloader
|
||||
class RGB2Lab(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs)
|
||||
return image_lab
|
||||
|
||||
|
||||
class Guassian_noise(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
|
||||
def __init__(self, noise_sigma=0.1):
|
||||
self.noise_sigma = noise_sigma
|
||||
|
||||
def __call__(self, inputs):
|
||||
h = inputs.shape[0]
|
||||
w = inputs.shape[1]
|
||||
noisy_image = inputs
|
||||
noise = np.random.randn(h, w) * self.noise_sigma
|
||||
noisy_image[:, :, 0] = inputs[:, :, 0] + noise
|
||||
|
||||
return noisy_image
|
||||
|
||||
|
||||
class Distortion(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
|
||||
def __init__(self, distortion_level=3, flip_probability=0.1):
|
||||
self.alpha_max = distortion_level
|
||||
self.flip_probability = flip_probability
|
||||
|
||||
def __call__(self, inputs):
|
||||
if np.random.rand() < self.flip_probability:
|
||||
inputs = inputs.transpose(Image.FLIP_LEFT_RIGHT)
|
||||
|
||||
inputs = np.array(inputs)
|
||||
alpha = np.random.rand() * self.alpha_max
|
||||
sigma = 50
|
||||
random_state = np.random.RandomState(None)
|
||||
shape = inputs.shape[0], inputs.shape[1]
|
||||
|
||||
dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
||||
dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
||||
|
||||
x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
|
||||
remap_image = cv2.remap(
|
||||
inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return Image.fromarray(remap_image)
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
|
||||
inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
|
||||
return inputs
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
outputs = F.to_mytensor(inputs)
|
||||
return outputs
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
"""
|
||||
Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al.
|
||||
-------------------------------------------------------------------------------------
|
||||
probability: The probability that the operation will be performed.
|
||||
sl: min erasing area
|
||||
sh: max erasing area
|
||||
r1: min aspect ratio
|
||||
mean: erasing value
|
||||
-------------------------------------------------------------------------------------
|
||||
"""
|
||||
|
||||
def __init__(self, probability=0.6, sl=0.05, sh=0.6):
|
||||
self.probability = probability
|
||||
self.sl = sl
|
||||
self.sh = sh
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
if random.uniform(0, 1) > self.probability:
|
||||
return Image.fromarray(img)
|
||||
|
||||
area = img.shape[0] * img.shape[1]
|
||||
h0 = img.shape[0]
|
||||
w0 = img.shape[1]
|
||||
channel = img.shape[2]
|
||||
|
||||
h = int(round(random.uniform(self.sl, self.sh) * h0))
|
||||
w = int(round(random.uniform(self.sl, self.sh) * w0))
|
||||
|
||||
if w < img.shape[1] and h < img.shape[0]:
|
||||
x1 = random.randint(0, img.shape[0] - h)
|
||||
y1 = random.randint(0, img.shape[1] - w)
|
||||
img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
|
||||
return Image.fromarray(img)
|
||||
|
||||
return Image.fromarray(img)
|
||||
|
||||
|
||||
class CenteredPad(object):
|
||||
"""
|
||||
pad the frame with black border,
|
||||
make square image for processing
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
self.image_size = image_size
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
width = np.size(img, 1)
|
||||
height = np.size(img, 0)
|
||||
old_size = [height, width]
|
||||
|
||||
ratio = float(self.image_size) / max(height, width)
|
||||
new_size = [int(x * ratio) for x in old_size]
|
||||
I_resize = imresize(img, new_size)
|
||||
width_new = np.size(I_resize, 1)
|
||||
height_new = np.size(I_resize, 0)
|
||||
|
||||
I_pad = np.zeros((self.image_size, self.image_size, 3))
|
||||
start_height = (self.image_size - new_size[0]) // 2
|
||||
start_width = (self.image_size - new_size[1]) // 2
|
||||
I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize
|
||||
|
||||
return Image.fromarray(I_pad.astype(np.uint8))
|
||||
@@ -0,0 +1,522 @@
|
||||
#!/usr/bin/python
|
||||
"""
|
||||
# ==============================
|
||||
# flowlib.py
|
||||
# library for optical flow processing
|
||||
# Author: Ruoteng Li
|
||||
# Date: 6th Aug 2016
|
||||
# ==============================
|
||||
"""
|
||||
import matplotlib.colors as cl
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import png
|
||||
from PIL import Image
|
||||
|
||||
UNKNOWN_FLOW_THRESH = 1e7
|
||||
SMALLFLOW = 0.0
|
||||
LARGEFLOW = 1e8
|
||||
|
||||
"""
|
||||
=============
|
||||
Flow Section
|
||||
=============
|
||||
"""
|
||||
|
||||
|
||||
def show_flow(filename):
|
||||
"""
|
||||
visualize optical flow map using matplotlib
|
||||
:param filename: optical flow file
|
||||
:return: None
|
||||
"""
|
||||
flow = read_flow(filename)
|
||||
img = flow_to_image(flow)
|
||||
plt.imshow(img)
|
||||
plt.show()
|
||||
|
||||
|
||||
def visualize_flow(flow, mode="Y"):
|
||||
"""
|
||||
this function visualize the input flow
|
||||
:param flow: input flow in array
|
||||
:param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color)
|
||||
:return: None
|
||||
"""
|
||||
if mode == "Y":
|
||||
# Ccbcr color wheel
|
||||
img = flow_to_image(flow)
|
||||
plt.imshow(img)
|
||||
plt.show()
|
||||
elif mode == "RGB":
|
||||
(h, w) = flow.shape[0:2]
|
||||
du = flow[:, :, 0]
|
||||
dv = flow[:, :, 1]
|
||||
valid = flow[:, :, 2]
|
||||
max_flow = max(np.max(du), np.max(dv))
|
||||
img = np.zeros((h, w, 3), dtype=np.float64)
|
||||
# angle layer
|
||||
img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi)
|
||||
# magnitude layer, normalized to 1
|
||||
img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow
|
||||
# phase layer
|
||||
img[:, :, 2] = 8 - img[:, :, 1]
|
||||
# clip to [0,1]
|
||||
small_idx = img[:, :, 0:3] < 0
|
||||
large_idx = img[:, :, 0:3] > 1
|
||||
img[small_idx] = 0
|
||||
img[large_idx] = 1
|
||||
# convert to rgb
|
||||
img = cl.hsv_to_rgb(img)
|
||||
# remove invalid point
|
||||
img[:, :, 0] = img[:, :, 0] * valid
|
||||
img[:, :, 1] = img[:, :, 1] * valid
|
||||
img[:, :, 2] = img[:, :, 2] * valid
|
||||
# show
|
||||
plt.imshow(img)
|
||||
plt.show()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def read_flow(filename):
|
||||
"""
|
||||
read optical flow from Middlebury .flo file
|
||||
:param filename: name of the flow file
|
||||
:return: optical flow data in matrix
|
||||
"""
|
||||
f = open(filename, "rb")
|
||||
try:
|
||||
magic = np.fromfile(f, np.float32, count=1)[0] # For Python3.x
|
||||
except:
|
||||
magic = np.fromfile(f, np.float32, count=1) # For Python2.x
|
||||
data2d = None
|
||||
|
||||
if 202021.25 != magic:
|
||||
print("Magic number incorrect. Invalid .flo file")
|
||||
else:
|
||||
w = np.fromfile(f, np.int32, count=1)[0]
|
||||
h = np.fromfile(f, np.int32, count=1)[0]
|
||||
# print("Reading %d x %d flo file" % (h, w))
|
||||
data2d = np.fromfile(f, np.float32, count=2 * w * h)
|
||||
# reshape data into 3D array (columns, rows, channels)
|
||||
data2d = np.resize(data2d, (h, w, 2))
|
||||
f.close()
|
||||
return data2d
|
||||
|
||||
|
||||
def read_flow_png(flow_file):
|
||||
"""
|
||||
Read optical flow from KITTI .png file
|
||||
:param flow_file: name of the flow file
|
||||
:return: optical flow data in matrix
|
||||
"""
|
||||
flow_object = png.Reader(filename=flow_file)
|
||||
flow_direct = flow_object.asDirect()
|
||||
flow_data = list(flow_direct[2])
|
||||
(w, h) = flow_direct[3]["size"]
|
||||
flow = np.zeros((h, w, 3), dtype=np.float64)
|
||||
for i in range(len(flow_data)):
|
||||
flow[i, :, 0] = flow_data[i][0::3]
|
||||
flow[i, :, 1] = flow_data[i][1::3]
|
||||
flow[i, :, 2] = flow_data[i][2::3]
|
||||
|
||||
invalid_idx = flow[:, :, 2] == 0
|
||||
flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0
|
||||
flow[invalid_idx, 0] = 0
|
||||
flow[invalid_idx, 1] = 0
|
||||
return flow
|
||||
|
||||
|
||||
def write_flow(flow, filename):
|
||||
"""
|
||||
write optical flow in Middlebury .flo format
|
||||
:param flow: optical flow map
|
||||
:param filename: optical flow file path to be saved
|
||||
:return: None
|
||||
"""
|
||||
f = open(filename, "wb")
|
||||
magic = np.array([202021.25], dtype=np.float32)
|
||||
(height, width) = flow.shape[0:2]
|
||||
w = np.array([width], dtype=np.int32)
|
||||
h = np.array([height], dtype=np.int32)
|
||||
magic.tofile(f)
|
||||
w.tofile(f)
|
||||
h.tofile(f)
|
||||
flow.tofile(f)
|
||||
f.close()
|
||||
|
||||
|
||||
def segment_flow(flow):
|
||||
h = flow.shape[0]
|
||||
w = flow.shape[1]
|
||||
u = flow[:, :, 0]
|
||||
v = flow[:, :, 1]
|
||||
|
||||
idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW)
|
||||
idx2 = abs(u) == SMALLFLOW
|
||||
class0 = (v == 0) & (u == 0)
|
||||
u[idx2] = 0.00001
|
||||
tan_value = v / u
|
||||
|
||||
class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0)
|
||||
class2 = (tan_value >= 1) & (u >= 0) & (v >= 0)
|
||||
class3 = (tan_value < -1) & (u <= 0) & (v >= 0)
|
||||
class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0)
|
||||
class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0)
|
||||
class7 = (tan_value < -1) & (u >= 0) & (v <= 0)
|
||||
class6 = (tan_value >= 1) & (u <= 0) & (v <= 0)
|
||||
class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0)
|
||||
|
||||
seg = np.zeros((h, w))
|
||||
|
||||
seg[class1] = 1
|
||||
seg[class2] = 2
|
||||
seg[class3] = 3
|
||||
seg[class4] = 4
|
||||
seg[class5] = 5
|
||||
seg[class6] = 6
|
||||
seg[class7] = 7
|
||||
seg[class8] = 8
|
||||
seg[class0] = 0
|
||||
seg[idx] = 0
|
||||
|
||||
return seg
|
||||
|
||||
|
||||
def flow_error(tu, tv, u, v):
|
||||
"""
|
||||
Calculate average end point error
|
||||
:param tu: ground-truth horizontal flow map
|
||||
:param tv: ground-truth vertical flow map
|
||||
:param u: estimated horizontal flow map
|
||||
:param v: estimated vertical flow map
|
||||
:return: End point error of the estimated flow
|
||||
"""
|
||||
smallflow = 0.0
|
||||
"""
|
||||
stu = tu[bord+1:end-bord,bord+1:end-bord]
|
||||
stv = tv[bord+1:end-bord,bord+1:end-bord]
|
||||
su = u[bord+1:end-bord,bord+1:end-bord]
|
||||
sv = v[bord+1:end-bord,bord+1:end-bord]
|
||||
"""
|
||||
stu = tu[:]
|
||||
stv = tv[:]
|
||||
su = u[:]
|
||||
sv = v[:]
|
||||
|
||||
idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH)
|
||||
stu[idxUnknow] = 0
|
||||
stv[idxUnknow] = 0
|
||||
su[idxUnknow] = 0
|
||||
sv[idxUnknow] = 0
|
||||
|
||||
ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)]
|
||||
index_su = su[ind2]
|
||||
index_sv = sv[ind2]
|
||||
an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1)
|
||||
un = index_su * an
|
||||
vn = index_sv * an
|
||||
|
||||
index_stu = stu[ind2]
|
||||
index_stv = stv[ind2]
|
||||
tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1)
|
||||
tun = index_stu * tn
|
||||
tvn = index_stv * tn
|
||||
|
||||
"""
|
||||
angle = un * tun + vn * tvn + (an * tn)
|
||||
index = [angle == 1.0]
|
||||
angle[index] = 0.999
|
||||
ang = np.arccos(angle)
|
||||
mang = np.mean(ang)
|
||||
mang = mang * 180 / np.pi
|
||||
"""
|
||||
|
||||
epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2)
|
||||
epe = epe[ind2]
|
||||
mepe = np.mean(epe)
|
||||
return mepe
|
||||
|
||||
|
||||
def flow_to_image(flow, display=False):
|
||||
"""
|
||||
Convert flow into middlebury color code image
|
||||
:param flow: optical flow map
|
||||
:return: optical flow image in middlebury color
|
||||
"""
|
||||
u = flow[:, :, 0]
|
||||
v = flow[:, :, 1]
|
||||
|
||||
maxu = -999.0
|
||||
maxv = -999.0
|
||||
minu = 999.0
|
||||
minv = 999.0
|
||||
|
||||
idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
|
||||
u[idxUnknow] = 0
|
||||
v[idxUnknow] = 0
|
||||
|
||||
maxu = max(maxu, np.max(u))
|
||||
minu = min(minu, np.min(u))
|
||||
|
||||
maxv = max(maxv, np.max(v))
|
||||
minv = min(minv, np.min(v))
|
||||
|
||||
rad = np.sqrt(u ** 2 + v ** 2)
|
||||
maxrad = max(-1, np.max(rad))
|
||||
|
||||
if display:
|
||||
print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv))
|
||||
|
||||
u = u / (maxrad + np.finfo(float).eps)
|
||||
v = v / (maxrad + np.finfo(float).eps)
|
||||
|
||||
img = compute_color(u, v)
|
||||
|
||||
idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
|
||||
img[idx] = 0
|
||||
|
||||
return np.uint8(img)
|
||||
|
||||
|
||||
def evaluate_flow_file(gt, pred):
|
||||
"""
|
||||
evaluate the estimated optical flow end point error according to ground truth provided
|
||||
:param gt: ground truth file path
|
||||
:param pred: estimated optical flow file path
|
||||
:return: end point error, float32
|
||||
"""
|
||||
# Read flow files and calculate the errors
|
||||
gt_flow = read_flow(gt) # ground truth flow
|
||||
eva_flow = read_flow(pred) # predicted flow
|
||||
# Calculate errors
|
||||
average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1])
|
||||
return average_pe
|
||||
|
||||
|
||||
def evaluate_flow(gt_flow, pred_flow):
|
||||
"""
|
||||
gt: ground-truth flow
|
||||
pred: estimated flow
|
||||
"""
|
||||
average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], pred_flow[:, :, 0], pred_flow[:, :, 1])
|
||||
return average_pe
|
||||
|
||||
|
||||
"""
|
||||
==============
|
||||
Disparity Section
|
||||
==============
|
||||
"""
|
||||
|
||||
|
||||
def read_disp_png(file_name):
|
||||
"""
|
||||
Read optical flow from KITTI .png file
|
||||
:param file_name: name of the flow file
|
||||
:return: optical flow data in matrix
|
||||
"""
|
||||
image_object = png.Reader(filename=file_name)
|
||||
image_direct = image_object.asDirect()
|
||||
image_data = list(image_direct[2])
|
||||
(w, h) = image_direct[3]["size"]
|
||||
channel = len(image_data[0]) / w
|
||||
flow = np.zeros((h, w, channel), dtype=np.uint16)
|
||||
for i in range(len(image_data)):
|
||||
for j in range(channel):
|
||||
flow[i, :, j] = image_data[i][j::channel]
|
||||
return flow[:, :, 0] / 256
|
||||
|
||||
|
||||
def disp_to_flowfile(disp, filename):
|
||||
"""
|
||||
Read KITTI disparity file in png format
|
||||
:param disp: disparity matrix
|
||||
:param filename: the flow file name to save
|
||||
:return: None
|
||||
"""
|
||||
f = open(filename, "wb")
|
||||
magic = np.array([202021.25], dtype=np.float32)
|
||||
(height, width) = disp.shape[0:2]
|
||||
w = np.array([width], dtype=np.int32)
|
||||
h = np.array([height], dtype=np.int32)
|
||||
empty_map = np.zeros((height, width), dtype=np.float32)
|
||||
data = np.dstack((disp, empty_map))
|
||||
magic.tofile(f)
|
||||
w.tofile(f)
|
||||
h.tofile(f)
|
||||
data.tofile(f)
|
||||
f.close()
|
||||
|
||||
|
||||
"""
|
||||
==============
|
||||
Image Section
|
||||
==============
|
||||
"""
|
||||
|
||||
|
||||
def read_image(filename):
|
||||
"""
|
||||
Read normal image of any format
|
||||
:param filename: name of the image file
|
||||
:return: image data in matrix uint8 type
|
||||
"""
|
||||
img = Image.open(filename)
|
||||
im = np.array(img)
|
||||
return im
|
||||
|
||||
|
||||
def warp_image(im, flow):
|
||||
"""
|
||||
Use optical flow to warp image to the next
|
||||
:param im: image to warp
|
||||
:param flow: optical flow
|
||||
:return: warped image
|
||||
"""
|
||||
from scipy import interpolate
|
||||
|
||||
image_height = im.shape[0]
|
||||
image_width = im.shape[1]
|
||||
flow_height = flow.shape[0]
|
||||
flow_width = flow.shape[1]
|
||||
n = image_height * image_width
|
||||
(iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width])
|
||||
(fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width])
|
||||
fx += flow[:, :, 0].astype(np.float)
|
||||
fy += flow[:, :, 1]
|
||||
mask = np.logical_or(fx < 0, fx > flow_width)
|
||||
mask = np.logical_or(mask, fy < 0)
|
||||
mask = np.logical_or(mask, fy > flow_height)
|
||||
fx = np.minimum(np.maximum(fx, 0), flow_width)
|
||||
fy = np.minimum(np.maximum(fy, 0), flow_height)
|
||||
points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1)
|
||||
xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1)
|
||||
warp = np.zeros((image_height, image_width, im.shape[2]))
|
||||
for i in range(im.shape[2]):
|
||||
channel = im[:, :, i]
|
||||
# plt.imshow(channel, cmap='gray')
|
||||
values = channel.reshape(n, 1)
|
||||
new_channel = interpolate.griddata(points, values, xi, method="cubic")
|
||||
new_channel = np.reshape(new_channel, [flow_height, flow_width])
|
||||
new_channel[mask] = 1
|
||||
warp[:, :, i] = new_channel.astype(np.uint8)
|
||||
|
||||
return warp.astype(np.uint8)
|
||||
|
||||
|
||||
"""
|
||||
==============
|
||||
Others
|
||||
==============
|
||||
"""
|
||||
|
||||
|
||||
def scale_image(image, new_range):
|
||||
"""
|
||||
Linearly scale the image into desired range
|
||||
:param image: input image
|
||||
:param new_range: the new range to be aligned
|
||||
:return: image normalized in new range
|
||||
"""
|
||||
min_val = np.min(image).astype(np.float32)
|
||||
max_val = np.max(image).astype(np.float32)
|
||||
min_val_new = np.array(min(new_range), dtype=np.float32)
|
||||
max_val_new = np.array(max(new_range), dtype=np.float32)
|
||||
scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new
|
||||
return scaled_image.astype(np.uint8)
|
||||
|
||||
|
||||
def compute_color(u, v):
|
||||
"""
|
||||
compute optical flow color map
|
||||
:param u: optical flow horizontal map
|
||||
:param v: optical flow vertical map
|
||||
:return: optical flow in color code
|
||||
"""
|
||||
[h, w] = u.shape
|
||||
img = np.zeros([h, w, 3])
|
||||
nanIdx = np.isnan(u) | np.isnan(v)
|
||||
u[nanIdx] = 0
|
||||
v[nanIdx] = 0
|
||||
|
||||
colorwheel = make_color_wheel()
|
||||
ncols = np.size(colorwheel, 0)
|
||||
|
||||
rad = np.sqrt(u ** 2 + v ** 2)
|
||||
|
||||
a = np.arctan2(-v, -u) / np.pi
|
||||
|
||||
fk = (a + 1) / 2 * (ncols - 1) + 1
|
||||
|
||||
k0 = np.floor(fk).astype(int)
|
||||
|
||||
k1 = k0 + 1
|
||||
k1[k1 == ncols + 1] = 1
|
||||
f = fk - k0
|
||||
|
||||
for i in range(0, np.size(colorwheel, 1)):
|
||||
tmp = colorwheel[:, i]
|
||||
col0 = tmp[k0 - 1] / 255
|
||||
col1 = tmp[k1 - 1] / 255
|
||||
col = (1 - f) * col0 + f * col1
|
||||
|
||||
idx = rad <= 1
|
||||
col[idx] = 1 - rad[idx] * (1 - col[idx])
|
||||
notidx = np.logical_not(idx)
|
||||
|
||||
col[notidx] *= 0.75
|
||||
img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def make_color_wheel():
|
||||
"""
|
||||
Generate color wheel according Middlebury color code
|
||||
:return: Color wheel
|
||||
"""
|
||||
RY = 15
|
||||
YG = 6
|
||||
GC = 4
|
||||
CB = 11
|
||||
BM = 13
|
||||
MR = 6
|
||||
|
||||
ncols = RY + YG + GC + CB + BM + MR
|
||||
|
||||
colorwheel = np.zeros([ncols, 3])
|
||||
|
||||
col = 0
|
||||
|
||||
# RY
|
||||
colorwheel[0:RY, 0] = 255
|
||||
colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
|
||||
col += RY
|
||||
|
||||
# YG
|
||||
colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG))
|
||||
colorwheel[col : col + YG, 1] = 255
|
||||
col += YG
|
||||
|
||||
# GC
|
||||
colorwheel[col : col + GC, 1] = 255
|
||||
colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))
|
||||
col += GC
|
||||
|
||||
# CB
|
||||
colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB))
|
||||
colorwheel[col : col + CB, 2] = 255
|
||||
col += CB
|
||||
|
||||
# BM
|
||||
colorwheel[col : col + BM, 2] = 255
|
||||
colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))
|
||||
col += +BM
|
||||
|
||||
# MR
|
||||
colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
|
||||
colorwheel[col : col + MR, 0] = 255
|
||||
|
||||
return colorwheel
|
||||
@@ -0,0 +1,42 @@
|
||||
import argparse
|
||||
import gc
|
||||
import glob
|
||||
import math
|
||||
import os
|
||||
import platform
|
||||
import random
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pdb import set_trace
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from skimage import color, io
|
||||
|
||||
video_name = "video.mp4"
|
||||
frame_folder = "../../results/seq1_new/"
|
||||
|
||||
images = [
|
||||
img for img in os.listdir(frame_folder) if img.endswith(".jpg") or img.endswith(".bmp") or img.endswith(".png")
|
||||
]
|
||||
images.sort()
|
||||
frame = cv2.imread(os.path.join(frame_folder, images[0]))
|
||||
height, width, layers = frame.shape
|
||||
video = cv2.VideoWriter("video.avi", cv2.VideoWriter_fourcc("X", "V", "I", "D"), 15, (width, height))
|
||||
# video = cv2.VideoWriter('video.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (width, height))
|
||||
for image in images:
|
||||
video.write(cv2.imread(os.path.join(frame_folder, image)))
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
video.release()
|
||||
|
||||
import imageio
|
||||
|
||||
frames = []
|
||||
for image in images:
|
||||
frames.append(imageio.imread(os.path.join(frame_folder, image)))
|
||||
imageio.mimsave("movie.gif", frames)
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
import numpy as np
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
|
||||
def hist_match_module(source, template):
|
||||
"""
|
||||
Adjust the pixel values of a grayscale image such that its histogram
|
||||
matches that of a target image
|
||||
|
||||
Arguments:
|
||||
-----------
|
||||
source: np.ndarray
|
||||
Image to transform; the histogram is computed over the flattened
|
||||
array
|
||||
template: np.ndarray
|
||||
Template image; can have different dimensions to source
|
||||
Returns:
|
||||
-----------
|
||||
matched: np.ndarray
|
||||
The transformed output image
|
||||
"""
|
||||
|
||||
oldshape = source.shape
|
||||
source = source.ravel()
|
||||
template = template.ravel()
|
||||
|
||||
# get the set of unique pixel values and their corresponding indices and
|
||||
# counts
|
||||
s_values, bin_idx, s_counts = np.unique(source, return_inverse=True, return_counts=True)
|
||||
t_values, t_counts = np.unique(template, return_counts=True)
|
||||
|
||||
# take the cumsum of the counts and normalize by the number of pixels to
|
||||
# get the empirical cumulative distribution functions for the source and
|
||||
# template images (maps pixel value --> quantile)
|
||||
s_quantiles = np.cumsum(s_counts).astype(np.float64)
|
||||
s_quantiles /= s_quantiles[-1]
|
||||
t_quantiles = np.cumsum(t_counts).astype(np.float64)
|
||||
t_quantiles /= t_quantiles[-1]
|
||||
|
||||
# interpolate linearly to find the pixel values in the template image
|
||||
# that correspond most closely to the quantiles in the source image
|
||||
interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)
|
||||
|
||||
return interp_t_values[bin_idx].reshape(oldshape)
|
||||
|
||||
|
||||
def hist_match(image_name1, image_name2, output_name):
|
||||
"""
|
||||
Adjust the pixel values of a grayscale image such that its histogram
|
||||
matches that of a target image
|
||||
|
||||
Arguments:
|
||||
-----------
|
||||
image name1, image name2
|
||||
Returns:
|
||||
-----------
|
||||
output dir + imagename
|
||||
"""
|
||||
image1 = plt.imread(image_name1)
|
||||
image2 = plt.imread(image_name2)
|
||||
# source = face()[:,:,0]
|
||||
# template = ascent()
|
||||
# matched = hist_match(source, template)
|
||||
|
||||
# def ecdf(x):
|
||||
# """convenience function for computing the empirical CDF"""
|
||||
# vals, counts = np.unique(x, return_counts=True)
|
||||
# ecdf = np.cumsum(counts).astype(np.float64)
|
||||
# ecdf /= ecdf[-1]
|
||||
# return vals, ecdf
|
||||
|
||||
# x1, y1 = ecdf(source.ravel())
|
||||
# x2, y2 = ecdf(template.ravel())
|
||||
# x3, y3 = ecdf(matched.ravel())
|
||||
|
||||
# fig = plt.figure()
|
||||
# gs = plt.GridSpec(2, 3)
|
||||
# ax1 = fig.add_subplot(gs[0, 0])
|
||||
# ax2 = fig.add_subplot(gs[0, 1], sharex=ax1, sharey=ax1)
|
||||
# ax3 = fig.add_subplot(gs[0, 2], sharex=ax1, sharey=ax1)
|
||||
# ax4 = fig.add_subplot(gs[1, :])
|
||||
# for aa in (ax1, ax2, ax3):
|
||||
# aa.set_axis_off()
|
||||
|
||||
# ax1.imshow(source, cmap=plt.cm.gray)
|
||||
# ax1.set_title('Source')
|
||||
# ax2.imshow(template, cmap=plt.cm.gray)
|
||||
# ax2.set_title('template')
|
||||
# ax3.imshow(matched, cmap=plt.cm.gray)
|
||||
# ax3.set_title('Matched')
|
||||
|
||||
# ax4.plot(x1, y1 * 100, '-r', lw=3, label='Source')
|
||||
# ax4.plot(x2, y2 * 100, '-k', lw=3, label='Template')
|
||||
# ax4.plot(x3, y3 * 100, '--r', lw=3, label='Matched')
|
||||
# ax4.set_xlim(x1[0], x1[-1])
|
||||
# ax4.set_xlabel('Pixel value')
|
||||
# ax4.set_ylabel('Cumulative %')
|
||||
# ax4.legend(loc=5)
|
||||
|
||||
# plt.show()
|
||||
# print('finished testing')
|
||||
Executable
+62
@@ -0,0 +1,62 @@
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
|
||||
vgg_cfg = {
|
||||
"A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
|
||||
"B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
|
||||
"D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"],
|
||||
"E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"],
|
||||
}
|
||||
|
||||
|
||||
def _vgg_make_conv_layer_names(cfg, batch_norm=False):
|
||||
layers = []
|
||||
block_idx = 1
|
||||
layer_idx = 1
|
||||
for v in cfg:
|
||||
if v == "M":
|
||||
layers += ["pool%d" % block_idx]
|
||||
block_idx += 1
|
||||
layer_idx = 1
|
||||
else:
|
||||
if batch_norm:
|
||||
layers += [
|
||||
"conv%d_%d" % (block_idx, layer_idx),
|
||||
"conv%d_%dnorm" % (block_idx, layer_idx),
|
||||
"relu%d_%d" % (block_idx, layer_idx),
|
||||
]
|
||||
else:
|
||||
layers += ["conv%d_%d" % (block_idx, layer_idx), "relu%d_%d" % (block_idx, layer_idx)]
|
||||
layer_idx += 1
|
||||
return layers
|
||||
|
||||
|
||||
def _vgg_make_fc_layer_names():
|
||||
layers = ["fc6", "relu6", "dropout6", "fc7", "relu7", "dropout7", "fc8"]
|
||||
return layers
|
||||
|
||||
|
||||
class VGG19Before(nn.Module):
|
||||
def __init__(self, layer_name, part="conv"):
|
||||
super(VGG19Before, self).__init__()
|
||||
assert part == "conv" or part == "fc", "invalid part name"
|
||||
vgg_19 = models.vgg19(pretrained=True)
|
||||
|
||||
if part == "conv":
|
||||
vgg19_conv_layers = _vgg_make_conv_layer_names(vgg_cfg["E"], batch_norm=False)
|
||||
idx = vgg19_conv_layers.index(layer_name)
|
||||
assert 0 <= idx < len(vgg19_conv_layers), "cannot find %s in part %s" % (layer_name, part)
|
||||
self.features = nn.Sequential(*list(vgg_19.features.children())[: (idx + 1)])
|
||||
else:
|
||||
vgg19_fc_layers = _vgg_make_fc_layer_names()
|
||||
idx = vgg19_fc_layers.index(layer_name)
|
||||
assert 0 <= idx < len(vgg19_fc_layers), "cannot find %s in part %s" % (layer_name, part)
|
||||
self.features = vgg_19.features
|
||||
self.classifier = nn.Sequential(*list(vgg_19.classifier.children())[: (idx + 1)])
|
||||
|
||||
def forward(self, x):
|
||||
x = self.features(x)
|
||||
if hasattr(self, "classifier"):
|
||||
x = self.classifier(x)
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
## Created by: Hang Zhang, Rutgers University, Email: zhang.hang@rutgers.edu
|
||||
## Modified by Thomas Wolf, HuggingFace Inc., Email: thomas@huggingface.co
|
||||
## Copyright (c) 2017-2018
|
||||
##
|
||||
## This source code is licensed under the MIT-style license found in the
|
||||
## LICENSE file in the root directory of this source tree
|
||||
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
"""Encoding Data Parallel"""
|
||||
import functools
|
||||
import threading
|
||||
|
||||
import torch
|
||||
import torch.cuda.comm as comm
|
||||
from torch.autograd import Function, Variable
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.nn.parallel._functions import Broadcast, ReduceAddCoalesced
|
||||
from torch.nn.parallel.data_parallel import DataParallel
|
||||
from torch.nn.parallel.parallel_apply import get_a_var
|
||||
from torch.nn.parallel.scatter_gather import gather
|
||||
|
||||
torch_ver = torch.__version__[:3]
|
||||
|
||||
__all__ = ["allreduce", "DataParallelModel", "DataParallelCriterion", "patch_replication_callback"]
|
||||
|
||||
|
||||
def allreduce(*inputs):
|
||||
"""Cross GPU all reduce autograd operation for calculate mean and
|
||||
variance in SyncBN.
|
||||
"""
|
||||
return AllReduce.apply(*inputs)
|
||||
|
||||
|
||||
class AllReduce(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, num_inputs, *inputs):
|
||||
ctx.num_inputs = num_inputs
|
||||
ctx.target_gpus = [inputs[i].get_device() for i in range(0, len(inputs), num_inputs)]
|
||||
inputs = [inputs[i : i + num_inputs] for i in range(0, len(inputs), num_inputs)]
|
||||
# sort before reduce sum
|
||||
inputs = sorted(inputs, key=lambda i: i[0].get_device())
|
||||
results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
|
||||
outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
|
||||
return tuple([t for tensors in outputs for t in tensors])
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *inputs):
|
||||
inputs = [i.data for i in inputs]
|
||||
inputs = [inputs[i : i + ctx.num_inputs] for i in range(0, len(inputs), ctx.num_inputs)]
|
||||
results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
|
||||
outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
|
||||
return (None,) + tuple([Variable(t) for tensors in outputs for t in tensors])
|
||||
|
||||
|
||||
class Reduce(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, *inputs):
|
||||
ctx.target_gpus = [inputs[i].get_device() for i in range(len(inputs))]
|
||||
inputs = sorted(inputs, key=lambda i: i.get_device())
|
||||
return comm.reduce_add(inputs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, gradOutput):
|
||||
return Broadcast.apply(ctx.target_gpus, gradOutput)
|
||||
|
||||
|
||||
class DistributedDataParallelModel(DistributedDataParallel):
|
||||
"""Implements data parallelism at the module level for the DistributedDataParallel module.
|
||||
This container parallelizes the application of the given module by
|
||||
splitting the input across the specified devices by chunking in the
|
||||
batch dimension.
|
||||
In the forward pass, the module is replicated on each device,
|
||||
and each replica handles a portion of the input. During the backwards pass,
|
||||
gradients from each replica are summed into the original module.
|
||||
Note that the outputs are not gathered, please use compatible
|
||||
:class:`encoding.parallel.DataParallelCriterion`.
|
||||
The batch size should be larger than the number of GPUs used. It should
|
||||
also be an integer multiple of the number of GPUs so that each chunk is
|
||||
the same size (so that each GPU processes the same number of samples).
|
||||
Args:
|
||||
module: module to be parallelized
|
||||
device_ids: CUDA devices (default: all devices)
|
||||
Reference:
|
||||
Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
|
||||
Amit Agrawal. “Context Encoding for Semantic Segmentation.
|
||||
*The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
|
||||
Example::
|
||||
>>> net = encoding.nn.DistributedDataParallelModel(model, device_ids=[0, 1, 2])
|
||||
>>> y = net(x)
|
||||
"""
|
||||
|
||||
def gather(self, outputs, output_device):
|
||||
return outputs
|
||||
|
||||
|
||||
class DataParallelModel(DataParallel):
|
||||
"""Implements data parallelism at the module level.
|
||||
|
||||
This container parallelizes the application of the given module by
|
||||
splitting the input across the specified devices by chunking in the
|
||||
batch dimension.
|
||||
In the forward pass, the module is replicated on each device,
|
||||
and each replica handles a portion of the input. During the backwards pass,
|
||||
gradients from each replica are summed into the original module.
|
||||
Note that the outputs are not gathered, please use compatible
|
||||
:class:`encoding.parallel.DataParallelCriterion`.
|
||||
|
||||
The batch size should be larger than the number of GPUs used. It should
|
||||
also be an integer multiple of the number of GPUs so that each chunk is
|
||||
the same size (so that each GPU processes the same number of samples).
|
||||
|
||||
Args:
|
||||
module: module to be parallelized
|
||||
device_ids: CUDA devices (default: all devices)
|
||||
|
||||
Reference:
|
||||
Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
|
||||
Amit Agrawal. “Context Encoding for Semantic Segmentation.
|
||||
*The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
|
||||
|
||||
Example::
|
||||
|
||||
>>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
|
||||
>>> y = net(x)
|
||||
"""
|
||||
|
||||
def gather(self, outputs, output_device):
|
||||
return outputs
|
||||
|
||||
def replicate(self, module, device_ids):
|
||||
modules = super(DataParallelModel, self).replicate(module, device_ids)
|
||||
execute_replication_callbacks(modules)
|
||||
return modules
|
||||
|
||||
|
||||
class DataParallelCriterion(DataParallel):
|
||||
"""
|
||||
Calculate loss in multiple-GPUs, which balance the memory usage.
|
||||
The targets are splitted across the specified devices by chunking in
|
||||
the batch dimension. Please use together with :class:`encoding.parallel.DataParallelModel`.
|
||||
|
||||
Reference:
|
||||
Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
|
||||
Amit Agrawal. “Context Encoding for Semantic Segmentation.
|
||||
*The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
|
||||
|
||||
Example::
|
||||
|
||||
>>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
|
||||
>>> criterion = encoding.nn.DataParallelCriterion(criterion, device_ids=[0, 1, 2])
|
||||
>>> y = net(x)
|
||||
>>> loss = criterion(y, target)
|
||||
"""
|
||||
|
||||
def forward(self, inputs, *targets, **kwargs):
|
||||
# input should be already scatterd
|
||||
# scattering the targets instead
|
||||
if not self.device_ids:
|
||||
return self.module(inputs, *targets, **kwargs)
|
||||
targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
|
||||
if len(self.device_ids) == 1:
|
||||
return self.module(inputs, *targets[0], **kwargs[0])
|
||||
replicas = self.replicate(self.module, self.device_ids[: len(inputs)])
|
||||
outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
|
||||
# return Reduce.apply(*outputs) / len(outputs)
|
||||
# return self.gather(outputs, self.output_device).mean()
|
||||
return self.gather(outputs, self.output_device)
|
||||
|
||||
|
||||
def _criterion_parallel_apply(modules, inputs, targets, kwargs_tup=None, devices=None):
|
||||
assert len(modules) == len(inputs)
|
||||
assert len(targets) == len(inputs)
|
||||
if kwargs_tup:
|
||||
assert len(modules) == len(kwargs_tup)
|
||||
else:
|
||||
kwargs_tup = ({},) * len(modules)
|
||||
if devices is not None:
|
||||
assert len(modules) == len(devices)
|
||||
else:
|
||||
devices = [None] * len(modules)
|
||||
|
||||
lock = threading.Lock()
|
||||
results = {}
|
||||
if torch_ver != "0.3":
|
||||
grad_enabled = torch.is_grad_enabled()
|
||||
|
||||
def _worker(i, module, input, target, kwargs, device=None):
|
||||
if torch_ver != "0.3":
|
||||
torch.set_grad_enabled(grad_enabled)
|
||||
if device is None:
|
||||
device = get_a_var(input).get_device()
|
||||
try:
|
||||
with torch.cuda.device(device):
|
||||
# this also avoids accidental slicing of `input` if it is a Tensor
|
||||
if not isinstance(input, (list, tuple)):
|
||||
input = (input,)
|
||||
if not isinstance(target, (list, tuple)):
|
||||
target = (target,)
|
||||
output = module(*(input + target), **kwargs)
|
||||
with lock:
|
||||
results[i] = output
|
||||
except Exception as e:
|
||||
with lock:
|
||||
results[i] = e
|
||||
|
||||
if len(modules) > 1:
|
||||
threads = [
|
||||
threading.Thread(target=_worker, args=(i, module, input, target, kwargs, device),)
|
||||
for i, (module, input, target, kwargs, device) in enumerate(
|
||||
zip(modules, inputs, targets, kwargs_tup, devices)
|
||||
)
|
||||
]
|
||||
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
else:
|
||||
_worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
|
||||
|
||||
outputs = []
|
||||
for i in range(len(inputs)):
|
||||
output = results[i]
|
||||
if isinstance(output, Exception):
|
||||
raise output
|
||||
outputs.append(output)
|
||||
return outputs
|
||||
|
||||
|
||||
###########################################################################
|
||||
# Adapted from Synchronized-BatchNorm-PyTorch.
|
||||
# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
|
||||
#
|
||||
class CallbackContext(object):
|
||||
pass
|
||||
|
||||
|
||||
def execute_replication_callbacks(modules):
|
||||
"""
|
||||
Execute an replication callback `__data_parallel_replicate__` on each module created
|
||||
by original replication.
|
||||
|
||||
The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
|
||||
|
||||
Note that, as all modules are isomorphism, we assign each sub-module with a context
|
||||
(shared among multiple copies of this module on different devices).
|
||||
Through this context, different copies can share some information.
|
||||
|
||||
We guarantee that the callback on the master copy (the first copy) will be called ahead
|
||||
of calling the callback of any slave copies.
|
||||
"""
|
||||
master_copy = modules[0]
|
||||
nr_modules = len(list(master_copy.modules()))
|
||||
ctxs = [CallbackContext() for _ in range(nr_modules)]
|
||||
|
||||
for i, module in enumerate(modules):
|
||||
for j, m in enumerate(module.modules()):
|
||||
if hasattr(m, "__data_parallel_replicate__"):
|
||||
m.__data_parallel_replicate__(ctxs[j], i)
|
||||
|
||||
|
||||
def patch_replication_callback(data_parallel):
|
||||
"""
|
||||
Monkey-patch an existing `DataParallel` object. Add the replication callback.
|
||||
Useful when you have customized `DataParallel` implementation.
|
||||
|
||||
Examples:
|
||||
> sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
|
||||
> sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
|
||||
> patch_replication_callback(sync_bn)
|
||||
# this is equivalent to
|
||||
> sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
|
||||
> sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
|
||||
"""
|
||||
|
||||
assert isinstance(data_parallel, DataParallel)
|
||||
|
||||
old_replicate = data_parallel.replicate
|
||||
|
||||
@functools.wraps(old_replicate)
|
||||
def new_replicate(module, device_ids):
|
||||
modules = old_replicate(module, device_ids)
|
||||
execute_replication_callbacks(modules)
|
||||
return modules
|
||||
|
||||
data_parallel.replicate = new_replicate
|
||||
@@ -0,0 +1,29 @@
|
||||
import threading
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class TBImageRecorder(threading.Thread):
|
||||
"""
|
||||
TBImageRecorder
|
||||
"""
|
||||
|
||||
def __init__(self, tb_writer, func, queue):
|
||||
super(TBImageRecorder, self).__init__()
|
||||
self._tb_writer = tb_writer
|
||||
# self._tb_writer = SummaryWriter()
|
||||
self._func = func
|
||||
self._queue = queue
|
||||
|
||||
def run(self):
|
||||
while True:
|
||||
msgs, iter_index = self._queue.get()
|
||||
if msgs:
|
||||
img_info = self._func(*msgs)
|
||||
|
||||
print("logging the images")
|
||||
for tag, images in img_info.items():
|
||||
if images is not None:
|
||||
self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index)
|
||||
else:
|
||||
break
|
||||
Executable
+67
@@ -0,0 +1,67 @@
|
||||
""""xyz2rgb(lab2xyz(lab, illuminant, observer))"""
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy import linalg
|
||||
from torch.autograd import Variable
|
||||
|
||||
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 = linalg.inv(xyz_from_rgb).T
|
||||
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
|
||||
# if neg_mask.any():
|
||||
# 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
|
||||
# if mask.any():
|
||||
# mask_xyz[mask] = torch.pow(xyz[mask], 3.)
|
||||
# if not mask.all():
|
||||
# mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.) / 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
|
||||
# if mask.any():
|
||||
# mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055
|
||||
# if not mask.all():
|
||||
# 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
|
||||
# if neg_mask.any():
|
||||
# mask_rgb[neg_mask] = 0
|
||||
# if large_mask.any():
|
||||
# mask_rgb[large_mask] = 1
|
||||
|
||||
return mask_rgb
|
||||
Executable
+88
@@ -0,0 +1,88 @@
|
||||
# Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514
|
||||
import numpy as np
|
||||
import scipy.misc
|
||||
import tensorflow as tf
|
||||
|
||||
try:
|
||||
from StringIO import StringIO # Python 2.7
|
||||
except ImportError:
|
||||
from io import BytesIO # Python 3.x
|
||||
|
||||
|
||||
class TBLogger(object):
|
||||
def __init__(self, log_dir):
|
||||
"""Create a summary writer logging to log_dir."""
|
||||
self.writer = tf.summary.FileWriter(log_dir)
|
||||
|
||||
def scalar_summary(self, tag, value, step):
|
||||
"""Log a scalar variable."""
|
||||
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
|
||||
self.writer.add_summary(summary, step)
|
||||
|
||||
def image_summaries(self, tag, images, step):
|
||||
"""Log a list of images."""
|
||||
|
||||
img_summaries = []
|
||||
for i, img in enumerate(images):
|
||||
# Write the image to a string
|
||||
try:
|
||||
s = StringIO()
|
||||
except:
|
||||
s = BytesIO()
|
||||
scipy.misc.toimage(img).save(s, format="png")
|
||||
|
||||
# Create an Image object
|
||||
img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=img.shape[0], width=img.shape[1])
|
||||
# Create a Summary value
|
||||
img_summaries.append(tf.Summary.Value(tag="%s/%d" % (tag, i), image=img_sum))
|
||||
|
||||
# Create and write Summary
|
||||
summary = tf.Summary(value=img_summaries)
|
||||
self.writer.add_summary(summary, step)
|
||||
|
||||
def image_summary(self, tag, image, step):
|
||||
"""Log a list of images."""
|
||||
|
||||
img_summaries = []
|
||||
|
||||
# Write the image to a string
|
||||
try:
|
||||
s = StringIO()
|
||||
except:
|
||||
s = BytesIO()
|
||||
|
||||
scipy.misc.toimage(image).save(s, format="png")
|
||||
|
||||
# Create an Image object
|
||||
img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=image.shape[0], width=image.shape[1])
|
||||
# Create and write Summary
|
||||
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, image=img_sum)])
|
||||
self.writer.add_summary(summary, step)
|
||||
|
||||
def histo_summary(self, tag, values, step, bins=1000):
|
||||
"""Log a histogram of the tensor of values."""
|
||||
|
||||
# Create a histogram using numpy
|
||||
counts, bin_edges = np.histogram(values, bins=bins)
|
||||
|
||||
# Fill the fields of the histogram proto
|
||||
hist = tf.HistogramProto()
|
||||
hist.min = float(np.min(values))
|
||||
hist.max = float(np.max(values))
|
||||
hist.num = int(np.prod(values.shape))
|
||||
hist.sum = float(np.sum(values))
|
||||
hist.sum_squares = float(np.sum(values ** 2))
|
||||
|
||||
# Drop the start of the first bin
|
||||
bin_edges = bin_edges[1:]
|
||||
|
||||
# Add bin edges and counts
|
||||
for edge in bin_edges:
|
||||
hist.bucket_limit.append(edge)
|
||||
for c in counts:
|
||||
hist.bucket.append(c)
|
||||
|
||||
# Create and write Summary
|
||||
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
|
||||
self.writer.add_summary(summary, step)
|
||||
self.writer.flush()
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
|
||||
import dominate
|
||||
from dominate.tags import *
|
||||
|
||||
|
||||
class HTML:
|
||||
def __init__(self, web_dir, title, reflesh=0):
|
||||
self.title = title
|
||||
self.web_dir = web_dir
|
||||
self.img_dir = os.path.join(self.web_dir, "images")
|
||||
if not os.path.exists(self.web_dir):
|
||||
os.makedirs(self.web_dir)
|
||||
if not os.path.exists(self.img_dir):
|
||||
os.makedirs(self.img_dir)
|
||||
# print(self.img_dir)
|
||||
|
||||
self.doc = dominate.document(title=title)
|
||||
if reflesh > 0:
|
||||
with self.doc.head:
|
||||
meta(http_equiv="reflesh", content=str(reflesh))
|
||||
|
||||
def get_image_dir(self):
|
||||
return self.img_dir
|
||||
|
||||
def add_header(self, str):
|
||||
with self.doc:
|
||||
h3(str)
|
||||
|
||||
def add_table(self, border=1):
|
||||
self.t = table(border=border, style="table-layout: fixed;")
|
||||
self.doc.add(self.t)
|
||||
|
||||
def add_images(self, ims, txts, links, width=400):
|
||||
self.add_table()
|
||||
with self.t:
|
||||
with tr():
|
||||
for im, txt, link in zip(ims, txts, links):
|
||||
with td(style="word-wrap: break-word;", halign="center", valign="top"):
|
||||
with p():
|
||||
with a(href=os.path.join("images", link)):
|
||||
img(style="width:%dpx" % width, src=os.path.join("images", im))
|
||||
br()
|
||||
p(txt)
|
||||
|
||||
def save(self):
|
||||
html_file = "%s/index.html" % self.web_dir
|
||||
f = open(html_file, "wt")
|
||||
f.write(self.doc.render())
|
||||
f.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
html = HTML("web/", "test_html")
|
||||
html.add_header("hello world")
|
||||
|
||||
ims = []
|
||||
txts = []
|
||||
links = []
|
||||
for n in range(4):
|
||||
ims.append("image_%d.jpg" % n)
|
||||
txts.append("text_%d" % n)
|
||||
links.append("image_%d.jpg" % n)
|
||||
html.add_images(ims, txts, links)
|
||||
html.add_images(ims, txts, links)
|
||||
html.save()
|
||||
@@ -0,0 +1,48 @@
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import skimage
|
||||
from skimage import color, feature, io, morphology
|
||||
from skimage.filters import prewitt, roberts, scharr, sobel
|
||||
from skimage.morphology import square
|
||||
|
||||
|
||||
def rgb2lms(rgb_image):
|
||||
lms_image = rgb_image
|
||||
r = rgb_image[:, :, 0]
|
||||
g = rgb_image[:, :, 1]
|
||||
b = rgb_image[:, :, 2]
|
||||
|
||||
l = 17.8824 * r + 43.5161 * g + 4.1194 * b
|
||||
m = 3.4557 * r + 27.1554 * g + 3.8671 * b
|
||||
s = 0.0300 * r + 0.1843 * g + 1.4671 * b
|
||||
lms_image = np.stack((l, m, s), axis=2)
|
||||
return l, m, s
|
||||
|
||||
|
||||
def image_normalize(rg_image):
|
||||
rg_image = rg_image - rg_image.min()
|
||||
rg_image = rg_image / (rg_image.max() - rg_image.min()) * 255
|
||||
return rg_image
|
||||
|
||||
|
||||
image_rgb = io.imread("ref5.jpg")
|
||||
plt.imshow(image_rgb, cmap="gray")
|
||||
L, M, S = rgb2lms(image_rgb)
|
||||
image_lab = color.rgb2lab(image_rgb)
|
||||
a_edge = sobel(image_lab[:, :, 1])
|
||||
b_edge = sobel(image_lab[:, :, 2])
|
||||
color_edge = np.sqrt(a_edge ** 2 + b_edge ** 2)
|
||||
luminance_edge = sobel(image_lab[:, :, 0])
|
||||
canny_edge_a = feature.canny(image_lab[:, :, 1], sigma=3, low_threshold=3)
|
||||
canny_edge_b = feature.canny(image_lab[:, :, 2], sigma=3, low_threshold=3)
|
||||
|
||||
plt.figure()
|
||||
# edge = morphology.opening(np.logical_and(color_edge > 8, luminance_edge > 4))
|
||||
edge = morphology.thin(
|
||||
morphology.closing((morphology.opening(np.logical_and(color_edge > 10, luminance_edge > 2))), square(1))
|
||||
)
|
||||
plt.imshow(edge, cmap="gray")
|
||||
|
||||
plt.figure()
|
||||
plt.imshow(np.logical_and(canny_edge_a, canny_edge_b), cmap="gray")
|
||||
plt.show()
|
||||
Executable
+422
@@ -0,0 +1,422 @@
|
||||
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
|
||||
@@ -0,0 +1,615 @@
|
||||
import math
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import lib.functional as F
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from scipy.misc import imread, imsave
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from skimage import color
|
||||
from skimage import transform as skimage_transform
|
||||
from skimage.draw import random_shapes
|
||||
from skimage.filters import gaussian
|
||||
from skimage.transform import resize, rotate
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
import time
|
||||
|
||||
import scipy.io as scio
|
||||
import skimage.io as io
|
||||
from numba import jit, u1, u2
|
||||
|
||||
|
||||
class RGB2Lab(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs)
|
||||
return image_lab
|
||||
|
||||
|
||||
class Guassian_noise(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
|
||||
def __init__(self, noise_sigma=0.1):
|
||||
self.noise_sigma = noise_sigma
|
||||
|
||||
def __call__(self, inputs):
|
||||
h = inputs.shape[0]
|
||||
w = inputs.shape[1]
|
||||
noisy_image = inputs
|
||||
noise = np.random.randn(h, w) * self.noise_sigma
|
||||
noisy_image[:, :, 0] = inputs[:, :, 0] + noise
|
||||
|
||||
return noisy_image
|
||||
|
||||
|
||||
class Distortion(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
|
||||
def __init__(self, distortion_level=3, flip_probability=0):
|
||||
self.alpha_max = distortion_level
|
||||
self.flip_probability = flip_probability
|
||||
|
||||
def __call__(self, inputs):
|
||||
if np.random.rand() < self.flip_probability:
|
||||
inputs = inputs.transpose(Image.FLIP_LEFT_RIGHT)
|
||||
|
||||
inputs = np.array(inputs)
|
||||
alpha = np.random.rand() * self.alpha_max
|
||||
sigma = 50
|
||||
random_state = np.random.RandomState(None)
|
||||
shape = inputs.shape[0], inputs.shape[1]
|
||||
|
||||
dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
||||
dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
||||
|
||||
x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
|
||||
remap_image = cv2.remap(
|
||||
inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return Image.fromarray(remap_image)
|
||||
|
||||
|
||||
class Distortion_with_flow(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def __call__(self, inputs, dx, dy):
|
||||
inputs = np.array(inputs)
|
||||
shape = inputs.shape[0], inputs.shape[1]
|
||||
inputs = np.array(inputs)
|
||||
# x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
|
||||
# remap_image = cv2.remap(inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR)
|
||||
# backward mapping
|
||||
|
||||
remap_image = forward_mapping(inputs, dy, dx, maxIter=3, precision=1e-3)
|
||||
# remap_image = cv2.remap(inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
|
||||
# remap_image = inputs
|
||||
|
||||
return Image.fromarray(remap_image)
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
|
||||
inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
|
||||
return inputs
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
outputs = F.to_mytensor(inputs)
|
||||
return outputs
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
"""
|
||||
Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al.
|
||||
-------------------------------------------------------------------------------------
|
||||
probability: The probability that the operation will be performed.
|
||||
sl: min erasing area
|
||||
sh: max erasing area
|
||||
r1: min aspect ratio
|
||||
mean: erasing value
|
||||
-------------------------------------------------------------------------------------
|
||||
"""
|
||||
|
||||
def __init__(self, probability=0.6, sl=0.05, sh=0.6):
|
||||
self.probability = probability
|
||||
self.sl = sl
|
||||
self.sh = sh
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
if random.uniform(0, 1) > self.probability:
|
||||
return Image.fromarray(img)
|
||||
|
||||
area = img.shape[0] * img.shape[1]
|
||||
h0 = img.shape[0]
|
||||
w0 = img.shape[1]
|
||||
channel = img.shape[2]
|
||||
|
||||
h = int(round(random.uniform(self.sl, self.sh) * h0))
|
||||
w = int(round(random.uniform(self.sl, self.sh) * w0))
|
||||
|
||||
if w < img.shape[1] and h < img.shape[0]:
|
||||
x1 = random.randint(0, img.shape[0] - h)
|
||||
y1 = random.randint(0, img.shape[1] - w)
|
||||
img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
|
||||
return Image.fromarray(img)
|
||||
|
||||
return Image.fromarray(img)
|
||||
|
||||
|
||||
class CenteredPad(object):
|
||||
"""
|
||||
pad the frame with black border,
|
||||
make square image for processing
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
self.image_size = image_size
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
width = np.size(img, 1)
|
||||
height = np.size(img, 0)
|
||||
old_size = [height, width]
|
||||
|
||||
ratio = float(self.image_size) / max(height, width)
|
||||
new_size = [int(x * ratio) for x in old_size]
|
||||
I_resize = resize(img, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_new = np.size(I_resize, 1)
|
||||
height_new = np.size(I_resize, 0)
|
||||
|
||||
I_pad = np.zeros((self.image_size, self.image_size, 3))
|
||||
start_height = (self.image_size - new_size[0]) // 2
|
||||
start_width = (self.image_size - new_size[1]) // 2
|
||||
I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize
|
||||
|
||||
return Image.fromarray(I_pad.astype(np.uint8))
|
||||
|
||||
|
||||
class centeredPad_with_height(object):
|
||||
"""
|
||||
pad the image according to the height
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
self.height = image_size[0]
|
||||
self.width = image_size[1]
|
||||
|
||||
def __call__(self, image):
|
||||
# pad the image to 16:9
|
||||
# pad height
|
||||
I = np.array(image)
|
||||
|
||||
# for padded input
|
||||
if I.shape[0] == I.shape[1] and I.shape[1] == self.width:
|
||||
return Image.fromarray(I.astype(np.uint8))
|
||||
|
||||
width0 = np.size(I, 1)
|
||||
height0 = np.size(I, 0)
|
||||
old_size = [height0, width0]
|
||||
height = self.height
|
||||
width = self.width
|
||||
|
||||
ratio = height / height0
|
||||
new_size = [int(x * ratio) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_new = np.size(I_resize, 1)
|
||||
height_new = np.size(I_resize, 0)
|
||||
|
||||
# if exceed the expected width
|
||||
if width_new > width:
|
||||
I_resize = I_resize[:, math.floor(width_new - width) // 2 : (math.floor(width_new - width) // 2 + width), :]
|
||||
width_new = np.size(I_resize, 1)
|
||||
height_new = np.size(I_resize, 0)
|
||||
|
||||
# lines: 56~200
|
||||
I_pad = np.zeros((width, width, 3))
|
||||
start_height = (width - height_new) // 2
|
||||
start_width = (width - width_new) // 2
|
||||
I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize
|
||||
|
||||
return Image.fromarray(I_pad.astype(np.uint8))
|
||||
|
||||
|
||||
class CenterPad(object):
|
||||
def __init__(self, image_size):
|
||||
self.height = image_size[0]
|
||||
self.width = image_size[1]
|
||||
|
||||
def __call__(self, image):
|
||||
# pad the image to 16:9
|
||||
# pad height
|
||||
I = np.array(image)
|
||||
|
||||
# for padded input
|
||||
height_old = np.size(I, 0)
|
||||
width_old = np.size(I, 1)
|
||||
old_size = [height_old, width_old]
|
||||
height = self.height
|
||||
width = self.width
|
||||
I_pad = np.zeros((height, width, np.size(I, 2)))
|
||||
|
||||
ratio = height / width
|
||||
if height_old / width_old == ratio:
|
||||
if height_old == height:
|
||||
return Image.fromarray(I.astype(np.uint8))
|
||||
else:
|
||||
new_size = [int(x * height / height_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
return Image.fromarray(I_resize.astype(np.uint8))
|
||||
|
||||
if height_old / width_old > ratio: # pad the width and crop
|
||||
new_size = [int(x * width / width_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_height = (height_resize - height) // 2
|
||||
I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :]
|
||||
else: # pad the height and crop
|
||||
new_size = [int(x * height / height_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_width = (width_resize - width) // 2
|
||||
I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :]
|
||||
|
||||
return Image.fromarray(I_pad.astype(np.uint8))
|
||||
|
||||
|
||||
class CenterPad_threshold(object):
|
||||
def __init__(self, image_size, threshold=3 / 4):
|
||||
self.height = image_size[0]
|
||||
self.width = image_size[1]
|
||||
self.threshold = threshold
|
||||
|
||||
def __call__(self, image):
|
||||
# pad the image to 16:9
|
||||
# pad height
|
||||
I = np.array(image)
|
||||
|
||||
# for padded input
|
||||
height_old = np.size(I, 0)
|
||||
width_old = np.size(I, 1)
|
||||
old_size = [height_old, width_old]
|
||||
height = self.height
|
||||
width = self.width
|
||||
I_pad = np.zeros((height, width, np.size(I, 2)))
|
||||
|
||||
ratio = height / width
|
||||
|
||||
if height_old / width_old == ratio:
|
||||
if height_old == height:
|
||||
return Image.fromarray(I.astype(np.uint8))
|
||||
else:
|
||||
new_size = [int(x * height / height_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
return Image.fromarray(I_resize.astype(np.uint8))
|
||||
|
||||
if height_old / width_old > self.threshold:
|
||||
width_new, height_new = width_old, int(width_old * self.threshold)
|
||||
height_margin = height_old - height_new
|
||||
height_crop_start = height_margin // 2
|
||||
I_crop = I[height_crop_start : (height_crop_start + height_new), :, :]
|
||||
I_resize = resize(
|
||||
I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True
|
||||
)
|
||||
|
||||
# debug:
|
||||
# height_new, width_new, channel = I_crop.shape
|
||||
# print(width_new / height_new)
|
||||
# import matplotlib.pyplot as plt
|
||||
# plt.imshow(I / 255)
|
||||
# plt.figure()
|
||||
# plt.imshow(I_resize / 255)
|
||||
return Image.fromarray(I_resize.astype(np.uint8))
|
||||
|
||||
if height_old / width_old > ratio: # pad the width and crop
|
||||
new_size = [int(x * width / width_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_height = (height_resize - height) // 2
|
||||
I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :]
|
||||
else: # pad the height and crop
|
||||
new_size = [int(x * height / height_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_width = (width_resize - width) // 2
|
||||
I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :]
|
||||
|
||||
return Image.fromarray(I_pad.astype(np.uint8))
|
||||
|
||||
|
||||
class CenterPadCrop_numpy(object):
|
||||
"""
|
||||
pad the image according to the height
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
self.height = image_size[0]
|
||||
self.width = image_size[1]
|
||||
|
||||
def __call__(self, image, threshold=3 / 4):
|
||||
# pad the image to 16:9
|
||||
# pad height
|
||||
I = np.array(image)
|
||||
# for padded input
|
||||
height_old = np.size(I, 0)
|
||||
width_old = np.size(I, 1)
|
||||
old_size = [height_old, width_old]
|
||||
height = self.height
|
||||
width = self.width
|
||||
padding_size = width
|
||||
if image.ndim == 2:
|
||||
I_pad = np.zeros((width, width))
|
||||
else:
|
||||
I_pad = np.zeros((width, width, I.shape[2]))
|
||||
|
||||
ratio = height / width
|
||||
if height_old / width_old == ratio:
|
||||
return I
|
||||
|
||||
if height_old / width_old > threshold:
|
||||
width_new, height_new = width_old, int(width_old * threshold)
|
||||
height_margin = height_old - height_new
|
||||
height_crop_start = height_margin // 2
|
||||
I_crop = I[height_start : (height_start + height_new), :]
|
||||
I_resize = resize(
|
||||
I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True
|
||||
)
|
||||
return I_resize
|
||||
|
||||
if height_old / width_old > ratio: # pad the width and crop
|
||||
new_size = [int(x * width / width_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_height = (height_resize - height) // 2
|
||||
start_height_block = (padding_size - height) // 2
|
||||
if image.ndim == 2:
|
||||
I_pad[start_height_block : (start_height_block + height), :] = I_resize[
|
||||
start_height : (start_height + height), :
|
||||
]
|
||||
else:
|
||||
I_pad[start_height_block : (start_height_block + height), :, :] = I_resize[
|
||||
start_height : (start_height + height), :, :
|
||||
]
|
||||
else: # pad the height and crop
|
||||
new_size = [int(x * height / height_old) for x in old_size]
|
||||
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
||||
width_resize = np.size(I_resize, 1)
|
||||
height_resize = np.size(I_resize, 0)
|
||||
start_width = (width_resize - width) // 2
|
||||
start_width_block = (padding_size - width) // 2
|
||||
if image.ndim == 2:
|
||||
I_pad[:, start_width_block : (start_width_block + width)] = I_resize[
|
||||
:, start_width : (start_width + width)
|
||||
]
|
||||
|
||||
else:
|
||||
I_pad[:, start_width_block : (start_width_block + width), :] = I_resize[
|
||||
:, start_width : (start_width + width), :
|
||||
]
|
||||
|
||||
crop_start_height = (I_pad.shape[0] - height) // 2
|
||||
crop_start_width = (I_pad.shape[1] - width) // 2
|
||||
|
||||
if image.ndim == 2:
|
||||
return I_pad[
|
||||
crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width)
|
||||
]
|
||||
else:
|
||||
return I_pad[
|
||||
crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width), :
|
||||
]
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision):
|
||||
# print('processing location', (xr, yr))
|
||||
#
|
||||
if abs(padu[yr, xr]) < precision and abs(padv[yr, xr]) < precision:
|
||||
return xr, yr
|
||||
|
||||
else:
|
||||
# Our initialize method in this paper, can see the overleaf for detail
|
||||
if (xr + 1) <= (W - 1):
|
||||
dif = padu[yr, xr + 1] - padu[yr, xr]
|
||||
u_next = padu[yr, xr] / (1 + dif)
|
||||
else:
|
||||
dif = padu[yr, xr] - padu[yr, xr - 1]
|
||||
u_next = padu[yr, xr] / (1 + dif)
|
||||
|
||||
if (yr + 1) <= (H - 1):
|
||||
dif = padv[yr + 1, xr] - padv[yr, xr]
|
||||
v_next = padv[yr, xr] / (1 + dif)
|
||||
else:
|
||||
dif = padv[yr, xr] - padv[yr - 1, xr]
|
||||
v_next = padv[yr, xr] / (1 + dif)
|
||||
|
||||
i = xr - u_next
|
||||
j = yr - v_next
|
||||
i_int = int(i)
|
||||
j_int = int(j)
|
||||
|
||||
# The same as traditional iterative search method
|
||||
for iter in range(maxIter):
|
||||
if 0 <= i <= (W - 1) and 0 <= j <= (H - 1):
|
||||
u11 = padu[j_int, i_int]
|
||||
v11 = padv[j_int, i_int]
|
||||
|
||||
u12 = padu[j_int, i_int + 1]
|
||||
v12 = padv[j_int, i_int + 1]
|
||||
|
||||
int1 = padu[j_int + 1, i_int]
|
||||
v21 = padv[j_int + 1, i_int]
|
||||
|
||||
int2 = padu[j_int + 1, i_int + 1]
|
||||
v22 = padv[j_int + 1, i_int + 1]
|
||||
|
||||
u = (
|
||||
u11 * (i_int + 1 - i) * (j_int + 1 - j)
|
||||
+ u12 * (i - i_int) * (j_int + 1 - j)
|
||||
+ int1 * (i_int + 1 - i) * (j - j_int)
|
||||
+ int2 * (i - i_int) * (j - j_int)
|
||||
)
|
||||
|
||||
v = (
|
||||
v11 * (i_int + 1 - i) * (j_int + 1 - j)
|
||||
+ v12 * (i - i_int) * (j_int + 1 - j)
|
||||
+ v21 * (i_int + 1 - i) * (j - j_int)
|
||||
+ v22 * (i - i_int) * (j - j_int)
|
||||
)
|
||||
|
||||
i_next = xr - u
|
||||
j_next = yr - v
|
||||
|
||||
if abs(i - i_next) < precision and abs(j - j_next) < precision:
|
||||
return i, j
|
||||
|
||||
i = i_next
|
||||
j = j_next
|
||||
|
||||
else:
|
||||
return i, j
|
||||
|
||||
# if the search doesn't converge within max iter, it will return the last iter result
|
||||
return i_next, j_next
|
||||
|
||||
|
||||
# Bilinear interpolation
|
||||
@jit(nopython=True)
|
||||
def biInterpolation(distorted, i, j):
|
||||
i = u2(i)
|
||||
j = u2(j)
|
||||
Q11 = distorted[j, i]
|
||||
Q12 = distorted[j, i + 1]
|
||||
Q21 = distorted[j + 1, i]
|
||||
Q22 = distorted[j + 1, i + 1]
|
||||
|
||||
pixel = u1(
|
||||
Q11 * (i + 1 - i) * (j + 1 - j)
|
||||
+ Q12 * (i - i) * (j + 1 - j)
|
||||
+ Q21 * (i + 1 - i) * (j - j)
|
||||
+ Q22 * (i - i) * (j - j)
|
||||
)
|
||||
return pixel
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter=5, precision=1e-2):
|
||||
for xr in range(W):
|
||||
for yr in range(H):
|
||||
# (xr, yr) is the point in result image, (i, j) is the search result in distorted image
|
||||
i, j = iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision)
|
||||
|
||||
# reflect the pixels outside the border
|
||||
if i > W - 1:
|
||||
i = 2 * W - 1 - i
|
||||
if i < 0:
|
||||
i = -i
|
||||
if j > H - 1:
|
||||
j = 2 * H - 1 - j
|
||||
if j < 0:
|
||||
j = -j
|
||||
|
||||
# Bilinear interpolation to get the pixel at (i, j) in distorted image
|
||||
resultImg[yr, xr, 0] = biInterpolation(distortImg[:, :, 0], i, j,)
|
||||
resultImg[yr, xr, 1] = biInterpolation(distortImg[:, :, 1], i, j,)
|
||||
resultImg[yr, xr, 2] = biInterpolation(distortImg[:, :, 2], i, j,)
|
||||
return None
|
||||
|
||||
|
||||
def forward_mapping(source_image, u, v, maxIter=5, precision=1e-2):
|
||||
"""
|
||||
warp the image according to the forward flow
|
||||
u: horizontal
|
||||
v: vertical
|
||||
"""
|
||||
H = source_image.shape[0]
|
||||
W = source_image.shape[1]
|
||||
|
||||
distortImg = np.array(np.zeros((H + 1, W + 1, 3)), dtype=np.uint8)
|
||||
distortImg[0:H, 0:W] = source_image[0:H, 0:W]
|
||||
distortImg[H, 0:W] = source_image[H - 1, 0:W]
|
||||
distortImg[0:H, W] = source_image[0:H, W - 1]
|
||||
distortImg[H, W] = source_image[H - 1, W - 1]
|
||||
|
||||
padu = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32)
|
||||
padu[0:H, 0:W] = u[0:H, 0:W]
|
||||
padu[H, 0:W] = u[H - 1, 0:W]
|
||||
padu[0:H, W] = u[0:H, W - 1]
|
||||
padu[H, W] = u[H - 1, W - 1]
|
||||
|
||||
padv = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32)
|
||||
padv[0:H, 0:W] = v[0:H, 0:W]
|
||||
padv[H, 0:W] = v[H - 1, 0:W]
|
||||
padv[0:H, W] = v[0:H, W - 1]
|
||||
padv[H, W] = v[H - 1, W - 1]
|
||||
|
||||
resultImg = np.array(np.zeros((H, W, 3)), dtype=np.uint8)
|
||||
iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter, precision)
|
||||
return resultImg
|
||||
|
||||
|
||||
def random_mask(H, W, mask_size=200):
|
||||
"""
|
||||
mask: ranges in [0,1]
|
||||
"""
|
||||
masked_image = np.zeros([H, W, 3])
|
||||
mask = random_shapes(
|
||||
[H, W],
|
||||
max_shapes=1,
|
||||
min_shapes=1,
|
||||
max_size=mask_size,
|
||||
min_size=mask_size / 2,
|
||||
multichannel=False,
|
||||
intensity_range=[0, 0],
|
||||
)[0]
|
||||
mask = np.stack((mask, mask, mask), axis=-1)
|
||||
random_state = np.random.RandomState(None)
|
||||
distortion_range = 50
|
||||
alpha = np.random.rand() * 6
|
||||
forward_dx = (
|
||||
gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000
|
||||
)
|
||||
forward_dy = (
|
||||
gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000
|
||||
)
|
||||
mask = forward_mapping(mask, forward_dy, forward_dx, maxIter=3, precision=1e-3) / 255
|
||||
mask = 1 - gaussian(mask, sigma=0, preserve_range=True, multichannel=False, anti_aliasing=True)
|
||||
mask = mask[:, :, 0]
|
||||
# masked_image = mask * image_foreground + (1 - mask) * image_background
|
||||
# masked_image = mask * image_foreground
|
||||
|
||||
# randomly transform the mask
|
||||
# rotate_degree = np.random.random() * 3.14
|
||||
# translation = [np.random.random() * H, np.random.random() * W]
|
||||
# tform = skimage_transform.SimilarityTransform(scale=1, rotation=rotate_degree, translation=translation)
|
||||
|
||||
# visualization
|
||||
# plt.imshow(masked_image / 255)
|
||||
# plt.show(block=True)
|
||||
return mask
|
||||
@@ -0,0 +1,14 @@
|
||||
from utils.util import to_np
|
||||
|
||||
|
||||
def histogram_logger(tb_writer, iter_index, net=None):
|
||||
if net is not None:
|
||||
for tag, value in net.named_parameters():
|
||||
tag = tag.replace(".", "/")
|
||||
tb_writer.add_histogram(tag, to_np(value), iter_index)
|
||||
tb_writer.add_histogram(tag + "/grad", to_np(value.grad), iter_index)
|
||||
|
||||
|
||||
def value_logger(tb_writer, iter_index, loss_info):
|
||||
for tag, value in loss_info.items():
|
||||
tb_writer.add_scalar(tag, value, iter_index)
|
||||
Executable
+136
@@ -0,0 +1,136 @@
|
||||
import collections
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def preprocess(img, scale_size=None):
|
||||
"""PILimg: RGB: HxWxC"""
|
||||
if scale_size is not None:
|
||||
scale_transforms = MaxScale(scale_size)
|
||||
img = scale_transforms(img)
|
||||
|
||||
prep_transforms = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to BGR
|
||||
transforms.Normalize(mean=[0.40760392, 0.45795686, 0.48501961], std=[1, 1, 1]), # subtract imagenet mean
|
||||
transforms.Lambda(lambda x: x.mul_(255)),
|
||||
]
|
||||
)
|
||||
img = prep_transforms(img)
|
||||
return img
|
||||
|
||||
|
||||
def deprocess(img):
|
||||
post_transforms_a = transforms.Compose(
|
||||
[
|
||||
transforms.Lambda(lambda x: x.mul_(1.0 / 255)),
|
||||
transforms.Normalize(mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1]), # add imagenet mean
|
||||
transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB
|
||||
]
|
||||
)
|
||||
post_transforms_b = transforms.Compose([transforms.ToPILImage()])
|
||||
img = post_transforms_a(img)
|
||||
img[img > 1] = 1
|
||||
img[img < 0] = 0
|
||||
img = post_transforms_b(img)
|
||||
return img
|
||||
|
||||
|
||||
def get_renamed_vgg():
|
||||
cache_file = "data/vgg19_conv.pth"
|
||||
vgg = models.vgg19().features
|
||||
renamed_vgg = nn.Sequential()
|
||||
part_idx, layer_idx = 1, 1
|
||||
for layer in list(vgg):
|
||||
if isinstance(layer, nn.Conv2d):
|
||||
name = "conv{}_{}".format(part_idx, layer_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
elif isinstance(layer, nn.ReLU):
|
||||
name = "relu{}_{}".format(part_idx, layer_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
layer_idx += 1
|
||||
elif isinstance(layer, nn.MaxPool2d):
|
||||
name = "pool{}".format(part_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
part_idx += 1
|
||||
layer_idx = 1
|
||||
renamed_vgg.load_state_dict(torch.load(cache_file))
|
||||
|
||||
return renamed_vgg
|
||||
|
||||
|
||||
def get_renamed_vgg_johnson():
|
||||
cache_file = "data/vgg19-d01eb7cb.pth"
|
||||
vgg = models.vgg19()
|
||||
model_dict = vgg.state_dict()
|
||||
pretrained_dict = torch.load(cache_file)
|
||||
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
|
||||
model_dict.update(pretrained_dict)
|
||||
vgg.load_state_dict(model_dict)
|
||||
vgg = vgg.features
|
||||
renamed_vgg = nn.Sequential()
|
||||
part_idx, layer_idx = 1, 1
|
||||
for layer in list(vgg):
|
||||
if isinstance(layer, nn.Conv2d):
|
||||
name = "conv{}_{}".format(part_idx, layer_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
elif isinstance(layer, nn.ReLU):
|
||||
name = "relu{}_{}".format(part_idx, layer_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
layer_idx += 1
|
||||
elif isinstance(layer, nn.MaxPool2d):
|
||||
name = "pool{}".format(part_idx)
|
||||
renamed_vgg.add_module(name, layer)
|
||||
part_idx += 1
|
||||
layer_idx = 1
|
||||
|
||||
return renamed_vgg
|
||||
|
||||
|
||||
class MaxScale(object):
|
||||
"""Rescale the input PIL.Image to the given size.
|
||||
|
||||
Args:
|
||||
size (sequence or int): Desired output size. If size is a sequence like
|
||||
(w, h), output size will be matched to this. If size is an int,
|
||||
smaller edge of the image will be matched to this number.
|
||||
i.e, if height > width, then image will be rescaled to
|
||||
(size * height / width, size)
|
||||
interpolation (int, optional): Desired interpolation. Default is
|
||||
``PIL.Image.BILINEAR``
|
||||
"""
|
||||
|
||||
def __init__(self, size, interpolation=Image.BILINEAR):
|
||||
assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
|
||||
self.size = size
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, img):
|
||||
"""
|
||||
Args:
|
||||
img (PIL.Image): Image to be scaled.
|
||||
|
||||
Returns:
|
||||
PIL.Image: Rescaled image.
|
||||
"""
|
||||
if isinstance(self.size, int):
|
||||
w, h = img.size
|
||||
if (w <= h == self.size) or (h <= w == self.size):
|
||||
return img
|
||||
if w < h:
|
||||
oh = self.size
|
||||
ow = int(self.size * w / h)
|
||||
return img.resize((ow, oh), self.interpolation)
|
||||
else:
|
||||
ow = self.size
|
||||
oh = int(self.size * h / w)
|
||||
return img.resize((ow, oh), self.interpolation)
|
||||
else:
|
||||
return img.resize(self.size, self.interpolation)
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
import glob
|
||||
import os
|
||||
import os.path as osp
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
import pdb
|
||||
|
||||
import lib.functional as F
|
||||
import matplotlib.pyplot as plt
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from skimage import color
|
||||
|
||||
from utils.util import (
|
||||
MovingAvg,
|
||||
batch_lab2rgb_transpose_mc,
|
||||
calc_ab_gradient,
|
||||
calc_cosine_dist_loss,
|
||||
calc_tv_loss,
|
||||
clean_tensorboard,
|
||||
feature_normalize,
|
||||
imshow,
|
||||
imshow_lab,
|
||||
mkdir_if_not,
|
||||
to_np,
|
||||
uncenter_l,
|
||||
vgg_preprocess,
|
||||
weighted_l1_loss,
|
||||
)
|
||||
|
||||
|
||||
def combo5_loader_nonoptimize(path):
|
||||
f = open(path, "rb")
|
||||
|
||||
# width, height
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
h = im_sz[0]
|
||||
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
w = im_sz[0]
|
||||
|
||||
# warp_ba_layer 4
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
d = f.read(im_sz[0])
|
||||
file_bytes = np.asarray(bytearray(d), dtype=np.uint8)
|
||||
img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
|
||||
img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB)
|
||||
warp_ba = Image.fromarray(img_data_ndarray)
|
||||
# warp_ba.save('/home/v-mingmh/Documents/gray/test_input/warp_ba.PNG')
|
||||
|
||||
# warp_aba_layer 4
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
d = f.read(im_sz[0])
|
||||
file_bytes = np.asarray(bytearray(d), dtype=np.uint8)
|
||||
img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
|
||||
img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB)
|
||||
warp_aba = Image.fromarray(img_data_ndarray)
|
||||
|
||||
# 5 layers: err_aba, err_ba, err_ab
|
||||
errs = []
|
||||
for l in range(5):
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
d = f.read(im_sz[0])
|
||||
file_bytes = np.asarray(bytearray(d), dtype=np.uint8)
|
||||
img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_GRAYSCALE)
|
||||
err_ba = Image.fromarray(img_data_ndarray)
|
||||
|
||||
# err_ba.save('/home/v-mingmh/Documents/gray/test_input/input_err_ba' + str(l) + '.PNG')
|
||||
|
||||
d = f.read(4)
|
||||
im_sz = struct.unpack("i", d)
|
||||
d = f.read(im_sz[0])
|
||||
file_bytes = np.asarray(bytearray(d), dtype=np.uint8)
|
||||
img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_GRAYSCALE)
|
||||
err_ab = Image.fromarray(img_data_ndarray)
|
||||
|
||||
# err_ab.save('/home/v-mingmh/Documents/gray/test_input/input_err_ab' + str(l) + '.PNG')
|
||||
|
||||
errs.append([err_ba, err_ab])
|
||||
|
||||
f.close()
|
||||
|
||||
return errs, warp_ba, warp_aba
|
||||
|
||||
|
||||
# combo_folder = '../../movie/seq1/combo_new/'
|
||||
combo_folder = "../../movie/seq2/combo_new"
|
||||
output_folder = "analogy_output2/"
|
||||
path, dirs, filenames = os.walk(combo_folder).__next__()
|
||||
|
||||
for filename in filenames:
|
||||
errs, warp_ba, warp_aba = combo5_loader_nonoptimize(combo_folder + "/" + filename)
|
||||
warp_ba.save(output_folder + filename.split(".")[0] + ".jpg")
|
||||
@@ -0,0 +1,34 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.autograd import Variable
|
||||
|
||||
|
||||
def get_grid(x):
|
||||
torchHorizontal = (
|
||||
torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3))
|
||||
)
|
||||
torchVertical = (
|
||||
torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3))
|
||||
)
|
||||
grid = torch.cat([torchHorizontal, torchVertical], 1).cuda()
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
class WarpingLayer(nn.Module):
|
||||
def __init__(self, device):
|
||||
super(WarpingLayer, self).__init__()
|
||||
self.device = device
|
||||
|
||||
def forward(self, x, flow):
|
||||
# WarpingLayer uses F.grid_sample, which expects normalized grid
|
||||
# we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code
|
||||
# so here we need to denormalize the flow
|
||||
flow_for_grip = torch.zeros_like(flow).cuda()
|
||||
flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0)
|
||||
flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0)
|
||||
|
||||
grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1)
|
||||
x_warp = F.grid_sample(x, grid)
|
||||
return x_warp
|
||||
Reference in New Issue
Block a user