From b87b7e7f258272116960ce8d90f18af3478e42bd Mon Sep 17 00:00:00 2001 From: zhangmozhe Date: Tue, 1 Dec 2020 14:12:10 +0800 Subject: [PATCH] refactor --- utils/tensor_lab2rgb.py | 47 ---------- utils/util_distortion.py | 184 ++++++++++++++++----------------------- utils/~frame2video.py | 26 ------ utils/~model_util.py | 61 ------------- 4 files changed, 76 insertions(+), 242 deletions(-) delete mode 100755 utils/tensor_lab2rgb.py delete mode 100644 utils/~frame2video.py delete mode 100755 utils/~model_util.py diff --git a/utils/tensor_lab2rgb.py b/utils/tensor_lab2rgb.py deleted file mode 100755 index c5b5533..0000000 --- a/utils/tensor_lab2rgb.py +++ /dev/null @@ -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 diff --git a/utils/util_distortion.py b/utils/util_distortion.py index 7e18b96..73dc45d 100644 --- a/utils/util_distortion.py +++ b/utils/util_distortion.py @@ -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 diff --git a/utils/~frame2video.py b/utils/~frame2video.py deleted file mode 100644 index 15cb7c1..0000000 --- a/utils/~frame2video.py +++ /dev/null @@ -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) diff --git a/utils/~model_util.py b/utils/~model_util.py deleted file mode 100755 index 1a824b9..0000000 --- a/utils/~model_util.py +++ /dev/null @@ -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