refactor
This commit is contained in:
zhangmozhe
+6
-6
@@ -17,7 +17,6 @@ class FeatVGG(nn.Module):
|
||||
def __init__(self, content_layers=["relu3_1"]):
|
||||
super(FeatVGG, self).__init__()
|
||||
self.content_layers = conv_to_relu(content_layers)
|
||||
# set up vgg19 model
|
||||
self.vgg19 = vgg_util.get_renamed_vgg()
|
||||
self.last_c_layer = self.content_layers[-1]
|
||||
is_last_content = False
|
||||
@@ -52,11 +51,12 @@ class VGGNet_multilayer(nn.Module):
|
||||
self.vgg = models.vgg19(pretrained=True).features
|
||||
|
||||
def forward(self, x):
|
||||
# x: rgb image
|
||||
# ranges in [0,1]
|
||||
# should be normalzied with mean = [0.485, 0.456, 0.406]
|
||||
# and variance = [0.229, 0.224, 0.225]
|
||||
"""Extract multiple convolutional feature maps."""
|
||||
"""Extract multiple convolutional feature maps.
|
||||
x: rgb image
|
||||
ranges in [0,1]
|
||||
should be normalzied with mean = [0.485, 0.456, 0.406]
|
||||
and variance = [0.229, 0.224, 0.225]
|
||||
"""
|
||||
features = []
|
||||
for name, layer in self.vgg._modules.items():
|
||||
x = layer(x)
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
from torch.autograd import Function, Variable
|
||||
|
||||
|
||||
class HardSigmoid(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
output = (0.2 * input) + 0.5
|
||||
output[output > 1] = 1
|
||||
output[output < 0] = 0
|
||||
ctx.save_for_backward(input)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
input = ctx.saved_variables
|
||||
grad_input = input.new(input.size())
|
||||
grad_input.data.fill_(0.2)
|
||||
grad_input[input <= -2.5] = 0
|
||||
grad_input[input >= 2.5] = 0
|
||||
return grad_input
|
||||
@@ -1,359 +0,0 @@
|
||||
from __future__ import division
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageEnhance, ImageOps
|
||||
|
||||
try:
|
||||
import accimage
|
||||
except ImportError:
|
||||
accimage = None
|
||||
import collections
|
||||
import numbers
|
||||
import types
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
from skimage import color
|
||||
|
||||
import lib.functional as F
|
||||
|
||||
__all__ = [
|
||||
"Compose",
|
||||
"Concatenate",
|
||||
"ToTensor",
|
||||
"Normalize",
|
||||
"Resize",
|
||||
"Scale",
|
||||
"CenterCrop",
|
||||
"Pad",
|
||||
"RandomCrop",
|
||||
"RandomHorizontalFlip",
|
||||
"RandomVerticalFlip",
|
||||
"RandomResizedCrop",
|
||||
"RandomSizedCrop",
|
||||
"FiveCrop",
|
||||
"TenCrop",
|
||||
"RGB2Lab",
|
||||
]
|
||||
|
||||
|
||||
def CustomFunc(inputs, func, *args, **kwargs):
|
||||
im_l = func(inputs[0], *args, **kwargs)
|
||||
im_ab = func(inputs[1], *args, **kwargs)
|
||||
warp_ba = func(inputs[2], *args, **kwargs)
|
||||
warp_aba = func(inputs[3], *args, **kwargs)
|
||||
im_gbl_ab = func(inputs[4], *args, **kwargs)
|
||||
bgr_mc_im = func(inputs[5], *args, **kwargs)
|
||||
|
||||
layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im]
|
||||
|
||||
for l in range(5):
|
||||
layer = inputs[6 + l]
|
||||
err_ba = func(layer[0], *args, **kwargs)
|
||||
err_ab = func(layer[1], *args, **kwargs)
|
||||
|
||||
layer_data.append([err_ba, err_ab])
|
||||
|
||||
return layer_data
|
||||
|
||||
|
||||
class Compose(object):
|
||||
"""Composes several transforms together.
|
||||
|
||||
Args:
|
||||
transforms (list of ``Transform`` objects): list of transforms to compose.
|
||||
|
||||
Example:
|
||||
>>> transforms.Compose([
|
||||
>>> transforms.CenterCrop(10),
|
||||
>>> transforms.ToTensor(),
|
||||
>>> ])
|
||||
"""
|
||||
|
||||
def __init__(self, transforms):
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, inputs):
|
||||
for t in self.transforms:
|
||||
inputs = t(inputs)
|
||||
return inputs
|
||||
|
||||
|
||||
class Concatenate(object):
|
||||
"""
|
||||
Input: [im_l, im_ab, inputs]
|
||||
inputs = [warp_ba_l, warp_ba_ab, warp_aba, err_pm, err_aba]
|
||||
|
||||
Output:[im_l, err_pm, warp_ba, warp_aba, im_ab, err_aba]
|
||||
"""
|
||||
|
||||
def __call__(self, inputs):
|
||||
im_l = inputs[0]
|
||||
im_ab = inputs[1]
|
||||
warp_ba = inputs[2]
|
||||
warp_aba = inputs[3]
|
||||
im_glb_ab = inputs[4]
|
||||
bgr_mc_im = inputs[5]
|
||||
bgr_mc_im = bgr_mc_im[[2, 1, 0], ...]
|
||||
|
||||
err_ba = []
|
||||
err_ab = []
|
||||
|
||||
for l in range(5):
|
||||
layer = inputs[6 + l]
|
||||
err_ba.append(layer[0])
|
||||
err_ab.append(layer[1])
|
||||
|
||||
cerr_ba = torch.cat(err_ba, 0)
|
||||
cerr_ab = torch.cat(err_ab, 0)
|
||||
|
||||
return (im_l, cerr_ba, warp_ba, warp_aba, im_glb_ab, bgr_mc_im, im_ab, cerr_ab)
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
"""Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
|
||||
|
||||
Converts a PIL Image or numpy.ndarray (H x W x C) in the range
|
||||
[0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].
|
||||
"""
|
||||
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
|
||||
|
||||
Returns:
|
||||
Tensor: Converted image.
|
||||
"""
|
||||
inputs = CustomFunc(inputs, F.to_mytensor)
|
||||
return inputs
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
"""Normalize an tensor image with mean and standard deviation.
|
||||
Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform
|
||||
will normalize each channel of the input ``torch.*Tensor`` i.e.
|
||||
``input[channel] = (input[channel] - mean[channel]) / std[channel]``
|
||||
|
||||
Args:
|
||||
mean (sequence): Sequence of means for each channel.
|
||||
std (sequence): Sequence of standard deviations for each channel.
|
||||
"""
|
||||
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
|
||||
|
||||
Returns:
|
||||
Tensor: Normalized Tensor image.
|
||||
"""
|
||||
|
||||
im_l = F.normalize(inputs[0], 50, 1) # [0, 100]
|
||||
im_ab = F.normalize(inputs[1], (0, 0), (1, 1)) # [-100, 100]
|
||||
|
||||
inputs[2][0:1, :, :] = F.normalize(inputs[2][0:1, :, :], 50, 1)
|
||||
inputs[2][1:3, :, :] = F.normalize(inputs[2][1:3, :, :], (0, 0), (1, 1))
|
||||
warp_ba = inputs[2]
|
||||
|
||||
inputs[3][0:1, :, :] = F.normalize(inputs[3][0:1, :, :], 50, 1)
|
||||
inputs[3][1:3, :, :] = F.normalize(inputs[3][1:3, :, :], (0, 0), (1, 1))
|
||||
warp_aba = inputs[3]
|
||||
|
||||
im_gbl_ab = F.normalize(inputs[4], (0, 0), (1, 1)) # [-100, 100]
|
||||
|
||||
bgr_mc_im = F.normalize(inputs[5], (123.68, 116.78, 103.938), (1, 1, 1))
|
||||
|
||||
layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im]
|
||||
|
||||
for l in range(5):
|
||||
layer = inputs[6 + l]
|
||||
err_ba = F.normalize(layer[0], 127, 2) # [0, 255]
|
||||
err_ab = F.normalize(layer[1], 127, 2) # [0, 255]
|
||||
layer_data.append([err_ba, err_ab])
|
||||
|
||||
return layer_data
|
||||
|
||||
|
||||
class Resize(object):
|
||||
"""Resize the input PIL Image to the given size.
|
||||
|
||||
Args:
|
||||
size (sequence or int): Desired output size. If size is a sequence like
|
||||
(h, w), 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, inputs):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be scaled.
|
||||
|
||||
Returns:
|
||||
PIL Image: Rescaled image.
|
||||
"""
|
||||
return CustomFunc(inputs, F.resize, self.size, self.interpolation)
|
||||
|
||||
|
||||
class RandomCrop(object):
|
||||
"""Crop the given PIL Image at a random location.
|
||||
|
||||
Args:
|
||||
size (sequence or int): Desired output size of the crop. If size is an
|
||||
int instead of sequence like (h, w), a square crop (size, size) is
|
||||
made.
|
||||
padding (int or sequence, optional): Optional padding on each border
|
||||
of the image. Default is 0, i.e no padding. If a sequence of length
|
||||
4 is provided, it is used to pad left, top, right, bottom borders
|
||||
respectively.
|
||||
"""
|
||||
|
||||
def __init__(self, size, padding=0):
|
||||
if isinstance(size, numbers.Number):
|
||||
self.size = (int(size), int(size))
|
||||
else:
|
||||
self.size = size
|
||||
self.padding = padding
|
||||
|
||||
@staticmethod
|
||||
def get_params(img, output_size):
|
||||
"""Get parameters for ``crop`` for a random crop.
|
||||
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped.
|
||||
output_size (tuple): Expected output size of the crop.
|
||||
|
||||
Returns:
|
||||
tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
|
||||
"""
|
||||
w, h = img.size
|
||||
th, tw = output_size
|
||||
if w == tw and h == th:
|
||||
return 0, 0, h, w
|
||||
|
||||
i = random.randint(0, h - th)
|
||||
j = random.randint(0, w - tw)
|
||||
return i, j, th, tw
|
||||
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped.
|
||||
|
||||
Returns:
|
||||
PIL Image: Cropped image.
|
||||
"""
|
||||
if self.padding > 0:
|
||||
inputs = CustomFunc(inputs, F.pad, self.padding)
|
||||
|
||||
i, j, h, w = self.get_params(inputs[0], self.size)
|
||||
return CustomFunc(inputs, F.crop, i, j, h, w)
|
||||
|
||||
|
||||
class CenterCrop(object):
|
||||
"""Crop the given PIL Image at a random location.
|
||||
|
||||
Args:
|
||||
size (sequence or int): Desired output size of the crop. If size is an
|
||||
int instead of sequence like (h, w), a square crop (size, size) is
|
||||
made.
|
||||
padding (int or sequence, optional): Optional padding on each border
|
||||
of the image. Default is 0, i.e no padding. If a sequence of length
|
||||
4 is provided, it is used to pad left, top, right, bottom borders
|
||||
respectively.
|
||||
"""
|
||||
|
||||
def __init__(self, size, padding=0):
|
||||
if isinstance(size, numbers.Number):
|
||||
self.size = (int(size), int(size))
|
||||
else:
|
||||
self.size = size
|
||||
self.padding = padding
|
||||
|
||||
@staticmethod
|
||||
def get_params(img, output_size):
|
||||
"""Get parameters for ``crop`` for a random crop.
|
||||
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped.
|
||||
output_size (tuple): Expected output size of the crop.
|
||||
|
||||
Returns:
|
||||
tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
|
||||
"""
|
||||
w, h = img.size
|
||||
th, tw = output_size
|
||||
if w == tw and h == th:
|
||||
return 0, 0, h, w
|
||||
|
||||
i = (h - th) // 2
|
||||
j = (w - tw) // 2
|
||||
return i, j, th, tw
|
||||
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be cropped.
|
||||
|
||||
Returns:
|
||||
PIL Image: Cropped image.
|
||||
"""
|
||||
if self.padding > 0:
|
||||
inputs = CustomFunc(inputs, F.pad, self.padding)
|
||||
|
||||
i, j, h, w = self.get_params(inputs[0], self.size)
|
||||
return CustomFunc(inputs, F.crop, i, j, h, w)
|
||||
|
||||
|
||||
class RandomHorizontalFlip(object):
|
||||
"""Horizontally flip the given PIL Image randomly with a probability of 0.5."""
|
||||
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be flipped.
|
||||
|
||||
Returns:
|
||||
PIL Image: Randomly flipped image.
|
||||
"""
|
||||
|
||||
if random.random() < 0.5:
|
||||
return CustomFunc(inputs, F.hflip)
|
||||
return inputs
|
||||
|
||||
|
||||
class RGB2Lab(object):
|
||||
def __call__(self, inputs):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be flipped.
|
||||
|
||||
Returns:
|
||||
PIL Image: Randomly flipped image.
|
||||
"""
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs[0])
|
||||
warp_ba_lab = color.rgb2lab(inputs[2])
|
||||
warp_aba_lab = color.rgb2lab(inputs[3])
|
||||
im_gbl_lab = color.rgb2lab(inputs[4])
|
||||
|
||||
inputs[0] = image_lab[:, :, :1] # l channel
|
||||
inputs[1] = image_lab[:, :, 1:] # ab channel
|
||||
inputs[2] = warp_ba_lab # lab channel
|
||||
inputs[3] = warp_aba_lab # lab channel
|
||||
inputs[4] = im_gbl_lab[:, :, 1:] # ab channel
|
||||
|
||||
return inputs
|
||||
-522
@@ -1,522 +0,0 @@
|
||||
#!/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
|
||||
+13
-14
@@ -208,24 +208,23 @@ def resize(img, size, interpolation=Image.BILINEAR):
|
||||
"""
|
||||
if not _is_pil_image(img):
|
||||
raise TypeError("img should be PIL Image. Got {}".format(type(img)))
|
||||
if not (isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)):
|
||||
if not isinstance(size, int) and (not isinstance(size, collections.Iterable) or len(size) != 2):
|
||||
raise TypeError("Got inappropriate size arg: {}".format(size))
|
||||
|
||||
if isinstance(size, int):
|
||||
w, h = img.size
|
||||
if (w <= h and w == size) or (h <= w and h == size):
|
||||
return img
|
||||
if w < h:
|
||||
ow = size
|
||||
oh = int(round(size * h / w))
|
||||
return img.resize((ow, oh), interpolation)
|
||||
else:
|
||||
oh = size
|
||||
ow = int(round(size * w / h))
|
||||
return img.resize((ow, oh), interpolation)
|
||||
else:
|
||||
if not isinstance(size, int):
|
||||
return img.resize(size[::-1], interpolation)
|
||||
|
||||
w, h = img.size
|
||||
if (w <= h and w == size) or (h <= w and h == size):
|
||||
return img
|
||||
if w < h:
|
||||
ow = size
|
||||
oh = int(round(size * h / w))
|
||||
else:
|
||||
oh = size
|
||||
ow = int(round(size * w / h))
|
||||
return img.resize((ow, oh), interpolation)
|
||||
|
||||
|
||||
def scale(*args, **kwargs):
|
||||
warnings.warn("The use of the transforms.Scale transform is deprecated, " + "please use transforms.Resize instead.")
|
||||
|
||||
@@ -1,179 +0,0 @@
|
||||
import os.path as osp
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from utils.util_distortion import CenterPadCrop_numpy, Distortion_with_flow, RGB2Lab, ToTensor, random_mask
|
||||
|
||||
import lib.functional as F
|
||||
from lib.RefPairDatasetWithGlobalOrigRGB_optimize import ToTensor, combo5_loader, image_loader, parse_images
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
|
||||
class ImageNetDataset(data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
epoch,
|
||||
image_size,
|
||||
with_bad=False,
|
||||
with_mid=False,
|
||||
transforms_imagenet=None,
|
||||
distortion_level=3,
|
||||
brightnessjitter=5,
|
||||
extra_reference_transform=None,
|
||||
real_reference_probability=1,
|
||||
mask_reference=False,
|
||||
no_reference_probability=0,
|
||||
):
|
||||
image_pairs = []
|
||||
curr_image_pairs = parse_images(data_root, with_bad, with_mid)
|
||||
image_pairs += curr_image_pairs
|
||||
print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs)))
|
||||
if len(image_pairs) == 0:
|
||||
raise RuntimeError("Found 0 image_a pairs in all the data_roots")
|
||||
|
||||
self.image_pairs = image_pairs
|
||||
self.epoch = epoch
|
||||
self.image_size = image_size
|
||||
self.real_len = len(self.image_pairs)
|
||||
self.image_pairs *= epoch
|
||||
|
||||
self.real_reference_probability = real_reference_probability
|
||||
self.no_reference_probability = no_reference_probability
|
||||
self.mask_reference = mask_reference
|
||||
|
||||
self.transforms_imagenet_raw = transforms_imagenet
|
||||
self.transforms_imagenet = transforms.Compose(transforms_imagenet)
|
||||
self.extra_reference_transform = transforms.Compose(extra_reference_transform)
|
||||
self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()])
|
||||
|
||||
self.distortion_level = distortion_level
|
||||
self.distortion_transform = Distortion_with_flow()
|
||||
self.brightnessjitter = brightnessjitter
|
||||
|
||||
def __getitem__(self, index):
|
||||
try:
|
||||
# print('imagenet loader')
|
||||
|
||||
image_id = 0
|
||||
pair_id = index
|
||||
|
||||
combo_path = None
|
||||
image_a_path = None
|
||||
image_b_path = None
|
||||
|
||||
image_names = ["", ""]
|
||||
dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id]
|
||||
sub_dir = osp.join(dir_root, cls_dir)
|
||||
if is_good >= 1:
|
||||
image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id])
|
||||
elif is_good == 0:
|
||||
image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id])
|
||||
else:
|
||||
image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id])
|
||||
|
||||
if np.random.random() > 0.5:
|
||||
image_a_path, image_b_path = image_b_path, image_a_path
|
||||
|
||||
I1 = image_loader(image_a_path)
|
||||
if np.random.random() < self.no_reference_probability:
|
||||
no_reference_flag = 1
|
||||
I_reference_fake = Image.fromarray(np.array(I1) * 0)
|
||||
I_reference_real = I_reference_fake
|
||||
else:
|
||||
no_reference_flag = 0
|
||||
I_reference_fake = I1
|
||||
I_reference_real = image_loader(image_b_path)
|
||||
|
||||
## generate the flow
|
||||
height, width = np.array(I1).shape[0], np.array(I1).shape[1]
|
||||
alpha = np.random.rand() * self.distortion_level
|
||||
distortion_range = 50
|
||||
random_state = np.random.RandomState(None)
|
||||
shape = self.image_size[0], self.image_size[1]
|
||||
# dx: flow on the vertical direction; dy: flow on the horizontal direction
|
||||
forward_dx = (
|
||||
gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
|
||||
* alpha
|
||||
* 1000
|
||||
)
|
||||
forward_dy = (
|
||||
gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
|
||||
* alpha
|
||||
* 1000
|
||||
)
|
||||
|
||||
# transform the input image for colorization
|
||||
I1 = self.transforms_imagenet(I1)
|
||||
|
||||
# DEBUG
|
||||
# from utils.util import batch_lab2rgb_transpose_mc
|
||||
# plt.imshow(image_loader(image_a_path))
|
||||
# plt.figure()
|
||||
# I1 = I1.unsqueeze(0)
|
||||
# plt.imshow(batch_lab2rgb_transpose_mc(I1[0:32, 0:1, :, :], I1[0:32, 1:3, :, :]))
|
||||
|
||||
# transform the fake reference
|
||||
I_reference_fake = self.extra_reference_transform(I_reference_fake)
|
||||
for transform in self.transforms_imagenet_raw:
|
||||
if type(transform) is RGB2Lab and self.extra_reference_transform is not None:
|
||||
I_reference_fake = self.distortion_transform(I_reference_fake, forward_dx, forward_dy)
|
||||
I_reference_fake = transform(I_reference_fake)
|
||||
I_reference_fake[0:1, :, :] = I_reference_fake[0:1, :, :] + torch.randn(1) * self.brightnessjitter
|
||||
|
||||
# transform the real reference
|
||||
I_reference_real = self.transforms_imagenet(I_reference_real)
|
||||
|
||||
# generate the masked image
|
||||
if self.mask_reference:
|
||||
mask = torch.Tensor(random_mask(self.image_size[0], self.image_size[1]))
|
||||
I_reference_fake[1:2, :, :] = (
|
||||
I_reference_fake[1:2, :, :] * (1 - mask) + I_reference_real[1:2, :, :] * mask
|
||||
)
|
||||
I_reference_fake[2:3, :, :] = (
|
||||
I_reference_fake[2:3, :, :] * (1 - mask) + I_reference_real[2:3, :, :] * mask
|
||||
)
|
||||
|
||||
# choose the reference
|
||||
if no_reference_flag:
|
||||
I_reference_output, fake_refernce_flag, no_reference_flag = (
|
||||
I_reference_real,
|
||||
torch.Tensor([0]),
|
||||
torch.Tensor([1]).unsqueeze(-1).unsqueeze(-1),
|
||||
)
|
||||
else:
|
||||
no_reference_flag = torch.Tensor([0]).unsqueeze(-1).unsqueeze(-1)
|
||||
I_reference_output, fake_refernce_flag = (
|
||||
(I_reference_real, torch.Tensor([0]))
|
||||
if np.random.random() < self.real_reference_probability
|
||||
else (I_reference_fake, torch.Tensor([1]))
|
||||
)
|
||||
|
||||
outputs = []
|
||||
outputs.append(I1)
|
||||
outputs.append(I_reference_output)
|
||||
outputs.append(fake_refernce_flag)
|
||||
outputs.append(no_reference_flag)
|
||||
|
||||
except Exception as e:
|
||||
if combo_path is not None:
|
||||
print("problem in ", combo_path)
|
||||
print("problem in, ", image_a_path)
|
||||
print(e)
|
||||
return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
|
||||
|
||||
return outputs
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_pairs)
|
||||
# return len(self.image_pairs) * 2
|
||||
@@ -1,82 +0,0 @@
|
||||
import collections
|
||||
import datetime
|
||||
import itertools
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
import re
|
||||
import string
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
from abc import abstractmethod
|
||||
from collections import Counter, Iterable, OrderedDict
|
||||
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from glob import glob, iglob
|
||||
from itertools import chain
|
||||
from operator import attrgetter, itemgetter
|
||||
|
||||
import bcolz
|
||||
import cv2
|
||||
import graphviz
|
||||
import IPython
|
||||
import matplotlib
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import PIL
|
||||
import scipy
|
||||
import seaborn as sns
|
||||
import sklearn
|
||||
import sklearn_pandas
|
||||
from IPython.lib.deepreload import reload as dreload
|
||||
from IPython.lib.display import FileLink
|
||||
from ipywidgets import fixed, interact, interactive, widgets
|
||||
from isoweek import Week
|
||||
from matplotlib import animation
|
||||
from matplotlib import pyplot as plt
|
||||
from matplotlib import rcParams
|
||||
from pandas_summary import DataFrameSummary
|
||||
from PIL import Image, ImageEnhance, ImageOps
|
||||
from sklearn import ensemble, metrics, preprocessing
|
||||
|
||||
matplotlib.rc("animation", html="html5")
|
||||
np.set_printoptions(precision=5, linewidth=110, suppress=True)
|
||||
|
||||
|
||||
def in_notebook():
|
||||
return "ipykernel" in sys.modules
|
||||
|
||||
|
||||
import tqdm as tq
|
||||
from tqdm import tnrange, tqdm_notebook
|
||||
|
||||
|
||||
def clear_tqdm():
|
||||
inst = getattr(tq.tqdm, "_instances", None)
|
||||
if not inst:
|
||||
return
|
||||
for i in range(len(inst)):
|
||||
inst.pop().close()
|
||||
|
||||
|
||||
if in_notebook():
|
||||
|
||||
def tqdm(*args, **kwargs):
|
||||
clear_tqdm()
|
||||
return tq.tqdm(*args, file=sys.stdout, **kwargs)
|
||||
|
||||
def trange(*args, **kwargs):
|
||||
clear_tqdm()
|
||||
return tq.trange(*args, file=sys.stdout, **kwargs)
|
||||
|
||||
|
||||
else:
|
||||
from tqdm import tqdm, trange
|
||||
|
||||
tnrange = trange
|
||||
tqdm_notebook = tqdm
|
||||
|
||||
@@ -1,390 +0,0 @@
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, "..")
|
||||
import os
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
import pdb
|
||||
import random
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import torch
|
||||
import torchvision
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from skimage import color
|
||||
from utils.flowlib import read_flow
|
||||
from utils.util_distortion import CenterPad, CenterPadCrop_numpy
|
||||
from utils.warping import WarpingLayer
|
||||
|
||||
import lib.functional as F
|
||||
|
||||
|
||||
class RGB2Lab(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs)
|
||||
return image_lab
|
||||
|
||||
|
||||
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) # permute channel and transform to tensor
|
||||
return outputs
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
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 CenterCrop(object):
|
||||
"""
|
||||
center crop the numpy array
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
# target size
|
||||
self.h0, self.w0 = image_size
|
||||
|
||||
def __call__(self, input_numpy):
|
||||
if input_numpy.ndim == 3:
|
||||
h, w, channel = input_numpy.shape
|
||||
output_numpy = np.zeros((self.h0, self.w0, channel))
|
||||
output_numpy = input_numpy[
|
||||
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, :
|
||||
]
|
||||
else:
|
||||
h, w = input_numpy.shape
|
||||
output_numpy = np.zeros((self.h0, self.w0))
|
||||
output_numpy = input_numpy[
|
||||
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0
|
||||
]
|
||||
return output_numpy
|
||||
|
||||
|
||||
def parse_images(data_root):
|
||||
image_pairs = []
|
||||
subdirs = sorted(os.listdir(data_root))
|
||||
for subdir in subdirs:
|
||||
path = os.path.join(data_root, subdir)
|
||||
if not os.path.isdir(path):
|
||||
continue
|
||||
|
||||
# parse_file = os.path.join(path, 'pairs_output.txt')
|
||||
parse_file = os.path.join(path, "pairs_output_new.txt")
|
||||
if os.path.exists(parse_file):
|
||||
f = open(parse_file, "r")
|
||||
lines = f.readlines()
|
||||
for line in lines:
|
||||
line = line.replace("\n", "")
|
||||
(
|
||||
image1_name,
|
||||
image2_name,
|
||||
reference_video_name,
|
||||
reference_video_name1,
|
||||
reference_name1,
|
||||
reference_name2,
|
||||
reference_name3,
|
||||
reference_name4,
|
||||
reference_name5,
|
||||
reference_gt1,
|
||||
reference_gt2,
|
||||
reference_gt3,
|
||||
) = line.split()
|
||||
image1_name = image1_name.split(".")[0]
|
||||
image2_name = image2_name.split(".")[0]
|
||||
reference_video_name = reference_video_name.split(".")[0]
|
||||
reference_video_name1 = reference_video_name1.split(".")[0]
|
||||
reference_name1 = reference_name1.split(".")[0]
|
||||
reference_name2 = reference_name2.split(".")[0]
|
||||
reference_name3 = reference_name3.split(".")[0]
|
||||
reference_name4 = reference_name4.split(".")[0]
|
||||
reference_name5 = reference_name5.split(".")[0]
|
||||
|
||||
reference_gt1 = reference_gt1.split(".")[0]
|
||||
reference_gt2 = reference_gt2.split(".")[0]
|
||||
reference_gt3 = reference_gt3.split(".")[0]
|
||||
|
||||
flow_forward_name = image1_name + "_forward"
|
||||
flow_backward_name = image1_name + "_backward"
|
||||
mask_name = image1_name + "_mask"
|
||||
|
||||
item = (
|
||||
image1_name + ".jpg",
|
||||
image2_name + ".jpg",
|
||||
reference_video_name + ".jpg",
|
||||
reference_name1 + ".JPEG",
|
||||
reference_name2 + ".JPEG",
|
||||
reference_name3 + ".JPEG",
|
||||
reference_name4 + ".JPEG",
|
||||
reference_name5 + ".JPEG",
|
||||
flow_forward_name + ".flo",
|
||||
flow_backward_name + ".flo",
|
||||
mask_name + ".pgm",
|
||||
reference_gt1 + ".jpg",
|
||||
reference_gt2 + ".jpg",
|
||||
reference_gt3 + ".jpg",
|
||||
path,
|
||||
)
|
||||
image_pairs.append(item)
|
||||
|
||||
f.close()
|
||||
|
||||
else:
|
||||
raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n"))
|
||||
|
||||
return image_pairs
|
||||
|
||||
|
||||
class VideosDataset(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
epoch,
|
||||
image_size,
|
||||
image_transform=None,
|
||||
use_google_reference=False,
|
||||
real_reference_probability=1,
|
||||
nonzero_placeholder_probability=0.5,
|
||||
):
|
||||
self.data_root = data_root
|
||||
self.image_transform = image_transform
|
||||
self.CenterPad = CenterPad(image_size)
|
||||
self.ToTensor = ToTensor()
|
||||
self.CenterCrop = CenterCrop(image_size)
|
||||
|
||||
assert len(self.data_root) > 0, "find no dataroot"
|
||||
self.epoch = epoch
|
||||
self.image_pairs = parse_images(self.data_root)
|
||||
self.real_len = len(self.image_pairs)
|
||||
print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len))
|
||||
self.image_pairs *= epoch
|
||||
self.use_google_reference = use_google_reference
|
||||
self.real_reference_probability = real_reference_probability
|
||||
self.nonzero_placeholder_probability = nonzero_placeholder_probability
|
||||
|
||||
def __getitem__(self, index):
|
||||
(
|
||||
image1_name,
|
||||
image2_name,
|
||||
reference_video_name,
|
||||
reference_name1,
|
||||
reference_name2,
|
||||
reference_name3,
|
||||
reference_name4,
|
||||
reference_name5,
|
||||
flow_forward_name,
|
||||
flow_backward_name,
|
||||
mask_name,
|
||||
reference_gt1,
|
||||
reference_gt2,
|
||||
reference_gt3,
|
||||
path,
|
||||
) = self.image_pairs[index]
|
||||
try:
|
||||
# print('video loader')
|
||||
I1 = Image.open(os.path.join(path, "input_pad", image1_name))
|
||||
I2 = Image.open(os.path.join(path, "input_pad", image2_name))
|
||||
|
||||
# randomly choose a reference frame from the video
|
||||
I_reference_video = Image.open(
|
||||
os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3]))
|
||||
)
|
||||
I_reference_video_real = Image.open(
|
||||
os.path.join(
|
||||
path,
|
||||
"reference",
|
||||
random.choice(
|
||||
[reference_name1, reference_name2, reference_name3, reference_name4, reference_name5]
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name)) # numpy
|
||||
flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name)) # numpy
|
||||
mask = Image.open(os.path.join(path, "mask", mask_name))
|
||||
|
||||
# binary mask
|
||||
mask = np.array(mask)
|
||||
mask[mask < 240] = 0
|
||||
mask[mask >= 240] = 1
|
||||
|
||||
# transform
|
||||
I1 = self.image_transform(I1)
|
||||
I2 = self.image_transform(I2)
|
||||
I_reference_video = self.image_transform(self.CenterPad(I_reference_video))
|
||||
I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real))
|
||||
flow_forward = self.ToTensor(self.CenterCrop(flow_forward))
|
||||
flow_backward = self.ToTensor(self.CenterCrop(flow_backward))
|
||||
mask = self.ToTensor(self.CenterCrop(mask))
|
||||
|
||||
if np.random.random() < self.real_reference_probability:
|
||||
I_reference_output = I_reference_video_real
|
||||
placeholder = torch.zeros_like(I1)
|
||||
self_ref_flag = torch.zeros_like(I1)
|
||||
else:
|
||||
I_reference_output = I_reference_video
|
||||
placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
|
||||
self_ref_flag = torch.ones_like(I1)
|
||||
|
||||
outputs = []
|
||||
outputs.append(I1)
|
||||
outputs.append(I2)
|
||||
outputs.append(I_reference_output)
|
||||
# outputs.append(I_reference_video_real)
|
||||
outputs.append(flow_forward)
|
||||
outputs.append(flow_backward)
|
||||
outputs.append(mask)
|
||||
outputs.append(placeholder)
|
||||
outputs.append(self_ref_flag)
|
||||
|
||||
except Exception as e:
|
||||
print("problem in, ", path)
|
||||
print(e)
|
||||
return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
|
||||
|
||||
return outputs
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_pairs)
|
||||
|
||||
|
||||
import torchvision.utils as vutils
|
||||
from torch.autograd import Variable
|
||||
|
||||
|
||||
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"
|
||||
|
||||
l_norm, ab_norm = 1.0, 1.0
|
||||
l_mean, ab_mean = 50.0, 0
|
||||
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
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
def test_videoloader():
|
||||
# dataset initI1lization
|
||||
# image size = 256
|
||||
train_dataset = VideosDataset(data_root="../../Dataset/video/video_pair", epoch=10)
|
||||
|
||||
data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
|
||||
|
||||
iter_instance = iter(data_loader)
|
||||
sample = iter_instance.next()
|
||||
I1_lab, I2_lab, I_reference_lab, flow_forward, flow_backward, mask = sample
|
||||
|
||||
# l: [-50,50], ab: [-128, 128]
|
||||
l_norm, ab_norm = 1.0, 1.0
|
||||
l_mean, ab_mean = 50.0, 0
|
||||
I1_l = I1_lab[:, 0:1, :, :]
|
||||
I1_ab = I1_lab[:, 1:3, :, :]
|
||||
I2_l = I2_lab[:, 0:1, :, :]
|
||||
I2_ab = I2_lab[:, 1:3, :, :]
|
||||
I_reference_l = I_reference_lab[:, 0:1, :, :]
|
||||
I_reference_ab = I_reference_lab[:, 1:3, :, :]
|
||||
|
||||
# test image samples
|
||||
image_A_grid = batch_lab2rgb_transpose_mc(I1_l, I1_ab, nrow=4).astype(np.uint8)
|
||||
image_B_grid = batch_lab2rgb_transpose_mc(I2_l, I2_ab, nrow=4).astype(np.uint8)
|
||||
image_ref_grid = batch_lab2rgb_transpose_mc(I_reference_l, I_reference_ab, nrow=4).astype(np.uint8)
|
||||
plt.imsave("A.png", image_A_grid)
|
||||
plt.imsave("B.png", image_B_grid)
|
||||
plt.figure()
|
||||
plt.imshow(image_A_grid)
|
||||
plt.title("image A")
|
||||
plt.figure()
|
||||
plt.imshow(image_B_grid)
|
||||
plt.title("image B")
|
||||
plt.figure()
|
||||
plt.imshow(image_ref_grid)
|
||||
plt.title("image ref")
|
||||
# plt.show()
|
||||
|
||||
# test warping
|
||||
print(I2_lab.shape)
|
||||
print(flow_forward.shape)
|
||||
warping_layer = WarpingLayer("gpu")
|
||||
I21_lab = warping_layer(I2_lab.cuda(), flow_forward.cuda())
|
||||
I21_lab = I21_lab * mask.cuda()
|
||||
|
||||
I21_l = I21_lab[:, 0:1, :, :]
|
||||
I21_ab = I21_lab[:, 1:3, :, :]
|
||||
image_warp_grid = batch_lab2rgb_transpose_mc(I21_l, I21_ab, nrow=4).astype(np.uint8)
|
||||
plt.imsave("BA.png", image_warp_grid)
|
||||
plt.figure()
|
||||
plt.imshow(image_warp_grid)
|
||||
plt.title("image warp")
|
||||
# plt.show()
|
||||
|
||||
# image difference
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(image_warp_grid - image_A_grid).astype(np.uint8))
|
||||
plt.title("warp-A")
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(image_B_grid - image_A_grid).astype(np.uint8))
|
||||
plt.title("B-A")
|
||||
# plt.show()
|
||||
plt.imsave("difference_warp.png", np.abs(image_warp_grid - image_A_grid).astype(np.uint8))
|
||||
plt.imsave("difference_BA.png", np.abs(image_B_grid - image_A_grid).astype(np.uint8))
|
||||
|
||||
return None
|
||||
@@ -1,203 +0,0 @@
|
||||
import os
|
||||
import os.path as osp
|
||||
import pdb
|
||||
import struct
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from scipy.misc import imread, imresize, imsave
|
||||
from scipy.ndimage import geometric_transform, map_coordinates
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from skimage import color
|
||||
from utils.util_distortion import (
|
||||
CenterPad,
|
||||
CenterPadCrop_numpy,
|
||||
Distortion,
|
||||
Distortion_with_flow,
|
||||
Normalize,
|
||||
RGB2Lab,
|
||||
ToTensor,
|
||||
)
|
||||
|
||||
import lib.functional as F
|
||||
from lib.RefPairDatasetWithGlobalOrigRGB_optimize import ToTensor, combo5_loader, image_loader, parse_images
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
|
||||
class VideosDataset_ImageNet(data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
epoch,
|
||||
image_size,
|
||||
with_bad=False,
|
||||
with_mid=False,
|
||||
transforms_imagenet=None,
|
||||
distortion_level=3,
|
||||
brightnessjitter=0,
|
||||
nonzero_placeholder_probability=0.5,
|
||||
extra_reference_transform=None,
|
||||
real_reference_probability=1,
|
||||
):
|
||||
image_pairs = []
|
||||
curr_image_pairs = parse_images(data_root, with_bad, with_mid)
|
||||
image_pairs += curr_image_pairs
|
||||
print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs)))
|
||||
if len(image_pairs) == 0:
|
||||
raise RuntimeError("Found 0 image_a pairs in all the data_roots")
|
||||
|
||||
self.image_pairs = image_pairs
|
||||
self.transforms_imagenet_raw = transforms_imagenet
|
||||
self.extra_reference_transform = transforms.Compose(extra_reference_transform)
|
||||
self.real_reference_probability = real_reference_probability
|
||||
self.transforms_imagenet = transforms.Compose(transforms_imagenet)
|
||||
self.epoch = epoch
|
||||
self.image_size = image_size
|
||||
self.real_len = len(self.image_pairs)
|
||||
self.image_pairs *= epoch
|
||||
self.distortion_level = distortion_level
|
||||
self.distortion_transform = Distortion_with_flow()
|
||||
self.brightnessjitter = brightnessjitter
|
||||
self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()])
|
||||
self.nonzero_placeholder_probability = nonzero_placeholder_probability
|
||||
self.ToTensor = ToTensor()
|
||||
self.Normalize = Normalize()
|
||||
|
||||
def __getitem__(self, index):
|
||||
try:
|
||||
# print('imagenet loader')
|
||||
|
||||
image_id = 0
|
||||
pair_id = index
|
||||
|
||||
combo_path = None
|
||||
image_a_path = None
|
||||
image_b_path = None
|
||||
|
||||
image_names = ["", ""]
|
||||
dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id]
|
||||
sub_dir = osp.join(dir_root, cls_dir)
|
||||
if is_good >= 1:
|
||||
image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id])
|
||||
elif is_good == 0:
|
||||
image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id])
|
||||
else:
|
||||
image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id])
|
||||
image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id])
|
||||
|
||||
if np.random.random() > 0.5:
|
||||
image_a_path, image_b_path = image_b_path, image_a_path
|
||||
|
||||
I1 = image_loader(image_a_path)
|
||||
I2 = I1
|
||||
I_reference_video = I1
|
||||
I_reference_video_real = image_loader(image_b_path)
|
||||
|
||||
## generate the flow
|
||||
height, width = np.array(I2).shape[0], np.array(I2).shape[1]
|
||||
alpha = np.random.rand() * self.distortion_level
|
||||
distortion_range = 50
|
||||
random_state = np.random.RandomState(None)
|
||||
shape = self.image_size[0], self.image_size[1]
|
||||
# dx: flow on the vertical direction; dy: flow on the horizontal direction
|
||||
forward_dx = (
|
||||
gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
|
||||
* alpha
|
||||
* 1000
|
||||
)
|
||||
forward_dy = (
|
||||
gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
|
||||
* alpha
|
||||
* 1000
|
||||
)
|
||||
|
||||
for transform in self.transforms_imagenet_raw:
|
||||
if type(transform) is RGB2Lab:
|
||||
I1_raw = I1
|
||||
I1 = transform(I1)
|
||||
else:
|
||||
I1 = transform(I1)
|
||||
|
||||
for transform in self.transforms_imagenet_raw:
|
||||
if type(transform) is RGB2Lab:
|
||||
I2 = self.distortion_transform(I2, forward_dx, forward_dy)
|
||||
I2_raw = I2
|
||||
I2 = transform(I2)
|
||||
I2[0:1, :, :] = I2[0:1, :, :] + torch.randn(1) * self.brightnessjitter
|
||||
|
||||
I_reference_video = self.extra_reference_transform(I_reference_video)
|
||||
for transform in self.transforms_imagenet_raw:
|
||||
# if type(transform) is RGB2Lab and self.extra_reference_transform is not None:
|
||||
# I_reference_video = self.extra_reference_transform(I_reference_video)
|
||||
I_reference_video = transform(I_reference_video)
|
||||
|
||||
I_reference_video_real = self.transforms_imagenet(I_reference_video_real)
|
||||
|
||||
flow_forward_raw = np.stack((forward_dy, forward_dx), axis=-1)
|
||||
flow_backward_raw = np.zeros_like(flow_forward_raw)
|
||||
flow_forward = self.flow_transform(flow_forward_raw)
|
||||
flow_backward = self.flow_transform(flow_backward_raw)
|
||||
|
||||
# update the mask for the pixels on the border
|
||||
grid_x, grid_y = np.meshgrid(np.arange(self.image_size[0]), np.arange(self.image_size[1]), indexing="ij")
|
||||
grid = np.stack((grid_y, grid_x), axis=-1)
|
||||
grid_warp = grid + flow_forward_raw
|
||||
location_y = grid_warp[:, :, 0].flatten()
|
||||
location_x = grid_warp[:, :, 1].flatten()
|
||||
I2_raw = np.array(I2_raw).astype(float)
|
||||
I21_r = map_coordinates(I2_raw[:, :, 0], np.stack((location_x, location_y)), cval=-1).reshape(
|
||||
(self.image_size[0], self.image_size[1])
|
||||
)
|
||||
I21_g = map_coordinates(I2_raw[:, :, 1], np.stack((location_x, location_y)), cval=-1).reshape(
|
||||
(self.image_size[0], self.image_size[1])
|
||||
)
|
||||
I21_b = map_coordinates(I2_raw[:, :, 2], np.stack((location_x, location_y)), cval=-1).reshape(
|
||||
(self.image_size[0], self.image_size[1])
|
||||
)
|
||||
I21_raw = np.stack((I21_r, I21_g, I21_b), axis=2)
|
||||
mask = np.ones((self.image_size[0], self.image_size[1]))
|
||||
mask[(I21_raw[:, :, 0] == -1) & (I21_raw[:, :, 1] == -1) & (I21_raw[:, :, 2] == -1)] = 0
|
||||
mask[abs(I21_raw - I1_raw).sum(axis=-1) > 50] = 0
|
||||
mask = self.ToTensor(mask)
|
||||
|
||||
if np.random.random() < self.real_reference_probability:
|
||||
I_reference_output = I_reference_video_real
|
||||
placeholder = torch.zeros_like(I1)
|
||||
self_ref_flag = torch.zeros_like(I1)
|
||||
else:
|
||||
I_reference_output = I_reference_video
|
||||
placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
|
||||
self_ref_flag = torch.ones_like(I1)
|
||||
|
||||
outputs = []
|
||||
outputs.append(I1)
|
||||
outputs.append(I2)
|
||||
outputs.append(I_reference_output)
|
||||
# outputs.append(I_reference_video_real)
|
||||
outputs.append(flow_forward)
|
||||
outputs.append(flow_backward)
|
||||
outputs.append(mask)
|
||||
outputs.append(placeholder)
|
||||
outputs.append(self_ref_flag)
|
||||
|
||||
except Exception as e:
|
||||
if combo_path is not None:
|
||||
print("problem in ", combo_path)
|
||||
print("problem in, ", image_a_path)
|
||||
print(e)
|
||||
return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
|
||||
|
||||
return outputs
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_pairs)
|
||||
# return len(self.image_pairs) * 2
|
||||
Reference in New Issue
Block a user