update training code
This commit is contained in:
zhangmozhe
@@ -0,0 +1,348 @@
|
||||
from __future__ import division
|
||||
|
||||
import collections
|
||||
import numbers
|
||||
import random
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
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.
|
||||
"""
|
||||
return CustomFunc(inputs, F.to_mytensor)
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,301 @@
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, "..")
|
||||
import os
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.utils as vutils
|
||||
from PIL import Image
|
||||
from skimage import color
|
||||
from torch.autograd import Variable
|
||||
from utils.flowlib import read_flow
|
||||
from utils.util_distortion import CenterPad
|
||||
|
||||
import lib.functional as F
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
|
||||
class RGB2Lab(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
return color.rgb2lab(inputs)
|
||||
|
||||
|
||||
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):
|
||||
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_new.txt")
|
||||
if os.path.exists(parse_file):
|
||||
with open(parse_file, "r") as f:
|
||||
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)
|
||||
|
||||
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:
|
||||
I1 = Image.open(os.path.join(path, "input_pad", image1_name))
|
||||
I2 = Image.open(os.path.join(path, "input_pad", image2_name))
|
||||
|
||||
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 = [
|
||||
I1,
|
||||
I2,
|
||||
I_reference_output,
|
||||
flow_forward,
|
||||
flow_backward,
|
||||
mask,
|
||||
placeholder,
|
||||
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)
|
||||
|
||||
|
||||
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")
|
||||
return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
||||
@@ -0,0 +1,290 @@
|
||||
import os
|
||||
import os.path as osp
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from scipy.ndimage import map_coordinates
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from utils.util_distortion import CenterPadCrop_numpy, Distortion_with_flow, Normalize, RGB2Lab, ToTensor
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
|
||||
def parse_images(dir, with_bad, with_mid):
|
||||
print("dir is: ", dir)
|
||||
dir = osp.expanduser(dir)
|
||||
|
||||
image_pairs = []
|
||||
|
||||
no_mid_cnt, no_bad_cnt = 0, 0
|
||||
for target in sorted(os.listdir(dir)):
|
||||
d = osp.join(dir, target)
|
||||
if not osp.isdir(d):
|
||||
continue
|
||||
|
||||
pair_file = osp.join(d, "pairs.txt")
|
||||
|
||||
if osp.exists(pair_file):
|
||||
with open(pair_file, "r") as f:
|
||||
lines = f.readlines()
|
||||
for line in lines:
|
||||
pair = line.strip().split(" ")
|
||||
name0, postfix = pair[0].split(".")
|
||||
name1, postfix = pair[1].split(".")
|
||||
if float(pair[2]) > 0:
|
||||
item0 = (dir, target, name0, name1, 2)
|
||||
item1 = (dir, target, name1, name0, 2)
|
||||
image_pairs.append(item0)
|
||||
image_pairs.append(item1)
|
||||
|
||||
else:
|
||||
raise (RuntimeError("Found no pair.txt in subfolders of: " + d + "\n"))
|
||||
|
||||
if with_mid:
|
||||
pair_file = osp.join(d, "pairs_mid.txt")
|
||||
|
||||
if osp.exists(pair_file):
|
||||
with open(pair_file, "r") as f:
|
||||
for line in f:
|
||||
pair = line.strip().split(" ")
|
||||
name0, postfix = pair[0].split(".")
|
||||
name1, postfix = pair[1].split(".")
|
||||
item0 = (dir, target, name0, name1, 0)
|
||||
item1 = (dir, target, name1, name0, 0)
|
||||
image_pairs.append(item0)
|
||||
image_pairs.append(item1)
|
||||
|
||||
else:
|
||||
no_mid_cnt += 1
|
||||
|
||||
if with_bad:
|
||||
pair_file = osp.join(d, "pairs_bad.txt")
|
||||
|
||||
if osp.exists(pair_file):
|
||||
with open(pair_file, "r") as f:
|
||||
for line in f:
|
||||
pair = line.strip().split(" ")
|
||||
name0, postfix = pair[0].split(".")
|
||||
name1, postfix = pair[1].split(".")
|
||||
item0 = (dir, target, name0, name1, -1)
|
||||
item1 = (dir, target, name1, name0, -1)
|
||||
image_pairs.append(item0)
|
||||
image_pairs.append(item1)
|
||||
|
||||
else:
|
||||
no_bad_cnt += 1
|
||||
|
||||
if no_bad_cnt > 0:
|
||||
print("find no pairs_bad.txt in %d folders for dir: %s" % (no_bad_cnt, dir))
|
||||
|
||||
if no_mid_cnt > 0:
|
||||
print("find no pairs_mid.txt in %d folders for dir: %s" % (no_mid_cnt, dir))
|
||||
|
||||
return image_pairs
|
||||
|
||||
|
||||
def pil_loader(path):
|
||||
with open(path, "rb") as f:
|
||||
with Image.open(f) as img:
|
||||
return img.convert("RGB")
|
||||
|
||||
|
||||
def image_loader(path):
|
||||
return pil_loader(path)
|
||||
|
||||
|
||||
def combo5_loader(path, real_w, real_h):
|
||||
with open(path, "rb") as f:
|
||||
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]
|
||||
|
||||
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)
|
||||
img_data_ndarray = img_data_ndarray * 0
|
||||
warp_aba = Image.fromarray(img_data_ndarray)
|
||||
|
||||
errs = []
|
||||
empty_array = np.zeros([np.array(warp_ba).shape[0], np.array(warp_ba).shape[1]])
|
||||
for _ in range(5):
|
||||
err_ba = Image.fromarray(empty_array)
|
||||
err_ab = Image.fromarray(empty_array)
|
||||
errs.append([err_ba, err_ab])
|
||||
return errs, warp_ba, warp_aba
|
||||
|
||||
|
||||
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 not image_pairs:
|
||||
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:
|
||||
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)
|
||||
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:
|
||||
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 = [
|
||||
I1,
|
||||
I2,
|
||||
I_reference_output,
|
||||
flow_forward,
|
||||
flow_backward,
|
||||
mask,
|
||||
placeholder,
|
||||
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)
|
||||
Reference in New Issue
Block a user