refactor
This commit is contained in:
zhangmozhe
@@ -1,47 +0,0 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
xyz_from_rgb = np.array(
|
||||
[[0.412453, 0.357580, 0.180423], [0.212671, 0.715160, 0.072169], [0.019334, 0.119193, 0.950227]]
|
||||
)
|
||||
rgb_from_xyz = np.array(
|
||||
[[3.24048134, -0.96925495, 0.05564664], [-1.53715152, 1.87599, -0.20404134], [-0.49853633, 0.04155593, 1.05731107]]
|
||||
)
|
||||
|
||||
|
||||
def tensor_lab2rgb(input):
|
||||
"""
|
||||
n * 3* h *w
|
||||
"""
|
||||
input_trans = input.transpose(1, 2).transpose(2, 3) # n * h * w * 3
|
||||
L, a, b = input_trans[:, :, :, 0:1], input_trans[:, :, :, 1:2], input_trans[:, :, :, 2:]
|
||||
y = (L + 16.0) / 116.0
|
||||
x = (a / 500.0) + y
|
||||
z = y - (b / 200.0)
|
||||
|
||||
neg_mask = z.data < 0
|
||||
z[neg_mask] = 0
|
||||
xyz = torch.cat((x, y, z), dim=3)
|
||||
|
||||
mask = xyz.data > 0.2068966
|
||||
mask_xyz = xyz.clone()
|
||||
mask_xyz[mask] = torch.pow(xyz[mask], 3.0)
|
||||
mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.0) / 7.787
|
||||
mask_xyz[:, :, :, 0] = mask_xyz[:, :, :, 0] * 0.95047
|
||||
mask_xyz[:, :, :, 2] = mask_xyz[:, :, :, 2] * 1.08883
|
||||
|
||||
rgb_trans = torch.mm(mask_xyz.view(-1, 3), torch.from_numpy(rgb_from_xyz).type_as(xyz)).view(
|
||||
input.size(0), input.size(2), input.size(3), 3
|
||||
)
|
||||
rgb = rgb_trans.transpose(2, 3).transpose(1, 2)
|
||||
|
||||
mask = rgb > 0.0031308
|
||||
mask_rgb = rgb.clone()
|
||||
mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055
|
||||
mask_rgb[~mask] = rgb[~mask] * 12.92
|
||||
|
||||
neg_mask = mask_rgb.data < 0
|
||||
large_mask = mask_rgb.data > 1
|
||||
mask_rgb[neg_mask] = 0
|
||||
mask_rgb[large_mask] = 1
|
||||
return mask_rgb
|
||||
+76
-108
@@ -3,23 +3,15 @@ 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
|
||||
from skimage.transform import resize
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
import time
|
||||
|
||||
import scipy.io as scio
|
||||
import skimage.io as io
|
||||
from numba import jit, u1, u2
|
||||
|
||||
|
||||
@@ -28,13 +20,11 @@ class RGB2Lab(object):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs)
|
||||
return image_lab
|
||||
return color.rgb2lab(inputs)
|
||||
|
||||
|
||||
class Guassian_noise(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
"""Elastic distortion"""
|
||||
|
||||
def __init__(self, noise_sigma=0.1):
|
||||
self.noise_sigma = noise_sigma
|
||||
@@ -50,8 +40,7 @@ class Guassian_noise(object):
|
||||
|
||||
|
||||
class Distortion(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
"""Elastic distortion"""
|
||||
|
||||
def __init__(self, distortion_level=3, flip_probability=0):
|
||||
self.alpha_max = distortion_level
|
||||
@@ -79,8 +68,7 @@ class Distortion(object):
|
||||
|
||||
|
||||
class Distortion_with_flow(object):
|
||||
"""Elastic distortion
|
||||
"""
|
||||
"""Elastic distortion"""
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
@@ -89,13 +77,7 @@ class Distortion_with_flow(object):
|
||||
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)
|
||||
|
||||
@@ -115,13 +97,12 @@ class ToTensor(object):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
outputs = F.to_mytensor(inputs)
|
||||
return outputs
|
||||
return F.to_mytensor(inputs)
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
"""
|
||||
Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al.
|
||||
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
|
||||
@@ -153,6 +134,7 @@ class RandomErasing(object):
|
||||
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)
|
||||
@@ -254,10 +236,9 @@ class CenterPad(object):
|
||||
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))
|
||||
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]
|
||||
@@ -301,10 +282,9 @@ class CenterPad_threshold(object):
|
||||
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))
|
||||
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)
|
||||
@@ -315,13 +295,6 @@ class CenterPad_threshold(object):
|
||||
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
|
||||
@@ -433,70 +406,64 @@ def iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision):
|
||||
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 (xr + 1) <= (W - 1):
|
||||
dif = padu[yr, xr + 1] - padu[yr, xr]
|
||||
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]
|
||||
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)
|
||||
|
||||
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)
|
||||
# The same as traditional iterative search method
|
||||
for _ in range(maxIter):
|
||||
if not 0 <= i <= (W - 1) or not 0 <= j <= (H - 1):
|
||||
return i, j
|
||||
|
||||
i = xr - u_next
|
||||
j = yr - v_next
|
||||
i_int = int(i)
|
||||
j_int = int(j)
|
||||
u11 = padu[j_int, i_int]
|
||||
v11 = padv[j_int, i_int]
|
||||
|
||||
# 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]
|
||||
|
||||
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]
|
||||
|
||||
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]
|
||||
|
||||
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)
|
||||
)
|
||||
|
||||
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)
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
i_next = xr - u
|
||||
j_next = yr - v
|
||||
if abs(i - i_next) < precision and abs(j - j_next) < precision:
|
||||
return i, j
|
||||
|
||||
if abs(i - i_next) < precision and abs(j - j_next) < precision:
|
||||
return i, j
|
||||
i = i_next
|
||||
j = j_next
|
||||
|
||||
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
|
||||
# if the search doesn't converge within max iter, it will return the last iter result
|
||||
return i_next, j_next
|
||||
|
||||
|
||||
# Bilinear interpolation
|
||||
@@ -509,13 +476,12 @@ def biInterpolation(distorted, i, j):
|
||||
Q21 = distorted[j + 1, i]
|
||||
Q22 = distorted[j + 1, i + 1]
|
||||
|
||||
pixel = u1(
|
||||
return 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)
|
||||
@@ -536,9 +502,21 @@ def iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter=5, precision=1e-
|
||||
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,)
|
||||
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
|
||||
|
||||
|
||||
@@ -601,15 +579,5 @@ def random_mask(H, W, mask_size=200):
|
||||
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
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
|
||||
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))
|
||||
|
||||
for image in images:
|
||||
video.write(cv2.imread(os.path.join(frame_folder, image)))
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
video.release()
|
||||
|
||||
import imageio
|
||||
|
||||
frames = [imageio.imread(os.path.join(frame_folder, image)) for image in images]
|
||||
|
||||
imageio.mimsave("movie.gif", frames)
|
||||
@@ -1,61 +0,0 @@
|
||||
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 in ["conv", "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
|
||||
Reference in New Issue
Block a user