init

lib/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.formatting.provider": "yapf"
+}

lib/FeatVGG.py

@@ -0,0 +1,65 @@
+import torch.nn as nn
+import utils.vgg_util as vgg_util
+from torchvision import models
+
+
+def conv_to_relu(layer_names):
+    out_layernames = []
+    for name in layer_names:
+        if name.startswith("conv"):
+            out_layernames.append("relu" + name[4:])
+        else:
+            out_layernames.append(name)
+    return out_layernames
+
+
+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
+        replace_layers, del_layers = [], []
+        for name, mod in self.vgg19.named_children():
+            if is_last_content:
+                del_layers.append(name)
+            else:
+                if name == self.last_c_layer:
+                    is_last_content = True
+
+        for name in del_layers:
+            delattr(self.vgg19, name)
+
+        # no need for gradweight vgg19
+        for param in self.vgg19.parameters():
+            param.requires_grad = False
+
+    def forward(self, input_img):
+        # input image is BGR image
+        # each channel ranges in [0,255]
+        # should be normalized with mean = [0.406*255, 0.456*255, 0.485*255] = [103,116,123]
+        # out = {}
+        return self.vgg19(input_img)
+
+
+class VGGNet_multilayer(nn.Module):
+    def __init__(self):
+        """Select conv1_1 ~ conv5_1 activation maps."""
+        super(VGGNet_multilayer, self).__init__()
+        self.select = ["0", "5", "10", "19", "28"]
+        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."""
+        features = []
+        for name, layer in self.vgg._modules.items():
+            x = layer(x)
+            if name in self.select:
+                features.append(x)
+        return features

lib/HardSigmoid.py

@@ -0,0 +1,20 @@
+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

lib/RefPairTransformsWithGlobalOrigRGB.py

@@ -0,0 +1,359 @@
+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

lib/TestTransforms.py

@@ -0,0 +1,353 @@
+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]
+    # layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im]
+    for l in range(5):
+        layer = inputs[4 + 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[4 + 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_ab, cerr_ab)
+        # 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]
+        layer_data = [im_l, im_ab, warp_ba, warp_aba]
+
+        for l in range(5):
+            layer = inputs[4 + 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)
+
+        if type(inputs) is list:
+            i, j, h, w = self.get_params(inputs[0], self.size)
+        else:
+            i, j, h, w = self.get_params(inputs, 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

lib/VGGFeatureLoss.py

@@ -0,0 +1,90 @@
+import torch.nn as nn
+import utils.vgg_util as vgg_util
+
+
+def conv_to_relu(layer_names):
+    out_layernames = []
+    for name in layer_names:
+        if name.startswith("conv"):
+            out_layernames.append("relu" + name[4:])
+        else:
+            out_layernames.append(name)
+    return out_layernames
+
+
+class VGGFeatureLoss(object):
+    def __init__(self, content_layers=["relu3_1"], content_weights=[1]):
+        self.content_layers = conv_to_relu(content_layers)
+        self.content_weights = content_weights
+        # set up vgg19 model
+        self.vgg19 = vgg_util.get_renamed_vgg()
+        self.last_c_layer = self.content_layers[-1]
+        is_last_content = False
+        replace_layers, del_layers = [], []
+        for name, mod in self.vgg19.named_children():
+            if is_last_content:
+                del_layers.append(name)
+            else:
+                if name == self.last_c_layer:
+                    is_last_content = True
+        print("del_layers", del_layers)
+        for name in del_layers:
+            delattr(self.vgg19, name)
+
+        # set up loss
+        self.loss_layers = self.content_layers
+        self.contentloss_fns = [nn.MSELoss()] * len(self.content_layers)
+        self.content_loss = None
+
+        # target
+        self.content_targets = None
+
+        # no need for gradweight vgg19
+        for param in self.vgg19.parameters():
+            param.requires_grad = False
+
+    def vgg_forward(self, input_img, layers, last_layer):
+        # out = {}
+        x = input_img
+        out = None
+        # print(last_layer)
+        assert len(layers) == 1 and layers[-1] == last_layer
+        for name, mod in self.vgg19.named_children():
+            # print('name is : ', name)
+            x = mod(x)
+            if name == last_layer:
+                out = x
+                break
+            # out[name] = mod(prev_input)
+            # prev_input = out[name]
+            # if name == last_layer:
+            #     break
+        return [out]
+        # return [out[key] for key in layers]
+
+    def cuda(self, gpu_ids):
+        if len(gpu_ids) > 1:
+            self.vgg19 = nn.DataParallel(self.vgg19.cuda(), gpu_ids)
+        else:
+            self.vgg19.cuda()
+        self.contentloss_fns = [loss_fn.cuda() for loss_fn in self.contentloss_fns]
+
+    def set_content_targets(self, content_img):
+        self.content_targets = [
+            A.detach() for A in self.vgg_forward(content_img, self.content_layers, self.last_c_layer)
+        ]
+
+    def forward(self, img):
+        assert self.content_targets is not None, "set up content targets first"
+        out = self.vgg_forward(img, self.content_layers, self.last_c_layer)
+        self.content_loss = sum(
+            [
+                self.content_weights[idx] * self.contentloss_fns[idx](out[idx], self.content_targets[idx])
+                for idx in range(len(self.content_layers))
+            ]
+        )
+        return self.content_loss
+
+    # def backward(self):
+    #     self.loss.backward()
+

lib/flowlib.py

@@ -0,0 +1,522 @@
+#!/usr/bin/python
+"""
+# ==============================
+# flowlib.py
+# library for optical flow processing
+# Author: Ruoteng Li
+# Date: 6th Aug 2016
+# ==============================
+"""
+import matplotlib.colors as cl
+import matplotlib.pyplot as plt
+import numpy as np
+import png
+from PIL import Image
+
+UNKNOWN_FLOW_THRESH = 1e7
+SMALLFLOW = 0.0
+LARGEFLOW = 1e8
+
+"""
+=============
+Flow Section
+=============
+"""
+
+
+def show_flow(filename):
+    """
+    visualize optical flow map using matplotlib
+    :param filename: optical flow file
+    :return: None
+    """
+    flow = read_flow(filename)
+    img = flow_to_image(flow)
+    plt.imshow(img)
+    plt.show()
+
+
+def visualize_flow(flow, mode="Y"):
+    """
+    this function visualize the input flow
+    :param flow: input flow in array
+    :param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color)
+    :return: None
+    """
+    if mode == "Y":
+        # Ccbcr color wheel
+        img = flow_to_image(flow)
+        plt.imshow(img)
+        plt.show()
+    elif mode == "RGB":
+        (h, w) = flow.shape[0:2]
+        du = flow[:, :, 0]
+        dv = flow[:, :, 1]
+        valid = flow[:, :, 2]
+        max_flow = max(np.max(du), np.max(dv))
+        img = np.zeros((h, w, 3), dtype=np.float64)
+        # angle layer
+        img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi)
+        # magnitude layer, normalized to 1
+        img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow
+        # phase layer
+        img[:, :, 2] = 8 - img[:, :, 1]
+        # clip to [0,1]
+        small_idx = img[:, :, 0:3] < 0
+        large_idx = img[:, :, 0:3] > 1
+        img[small_idx] = 0
+        img[large_idx] = 1
+        # convert to rgb
+        img = cl.hsv_to_rgb(img)
+        # remove invalid point
+        img[:, :, 0] = img[:, :, 0] * valid
+        img[:, :, 1] = img[:, :, 1] * valid
+        img[:, :, 2] = img[:, :, 2] * valid
+        # show
+        plt.imshow(img)
+        plt.show()
+
+    return None
+
+
+def read_flow(filename):
+    """
+    read optical flow from Middlebury .flo file
+    :param filename: name of the flow file
+    :return: optical flow data in matrix
+    """
+    f = open(filename, "rb")
+    try:
+        magic = np.fromfile(f, np.float32, count=1)[0]  # For Python3.x
+    except:
+        magic = np.fromfile(f, np.float32, count=1)  # For Python2.x
+    data2d = None
+
+    if 202021.25 != magic:
+        print("Magic number incorrect. Invalid .flo file")
+    else:
+        w = np.fromfile(f, np.int32, count=1)[0]
+        h = np.fromfile(f, np.int32, count=1)[0]
+        # print("Reading %d x %d flo file" % (h, w))
+        data2d = np.fromfile(f, np.float32, count=2 * w * h)
+        # reshape data into 3D array (columns, rows, channels)
+        data2d = np.resize(data2d, (h, w, 2))
+    f.close()
+    return data2d
+
+
+def read_flow_png(flow_file):
+    """
+    Read optical flow from KITTI .png file
+    :param flow_file: name of the flow file
+    :return: optical flow data in matrix
+    """
+    flow_object = png.Reader(filename=flow_file)
+    flow_direct = flow_object.asDirect()
+    flow_data = list(flow_direct[2])
+    (w, h) = flow_direct[3]["size"]
+    flow = np.zeros((h, w, 3), dtype=np.float64)
+    for i in range(len(flow_data)):
+        flow[i, :, 0] = flow_data[i][0::3]
+        flow[i, :, 1] = flow_data[i][1::3]
+        flow[i, :, 2] = flow_data[i][2::3]
+
+    invalid_idx = flow[:, :, 2] == 0
+    flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0
+    flow[invalid_idx, 0] = 0
+    flow[invalid_idx, 1] = 0
+    return flow
+
+
+def write_flow(flow, filename):
+    """
+    write optical flow in Middlebury .flo format
+    :param flow: optical flow map
+    :param filename: optical flow file path to be saved
+    :return: None
+    """
+    f = open(filename, "wb")
+    magic = np.array([202021.25], dtype=np.float32)
+    (height, width) = flow.shape[0:2]
+    w = np.array([width], dtype=np.int32)
+    h = np.array([height], dtype=np.int32)
+    magic.tofile(f)
+    w.tofile(f)
+    h.tofile(f)
+    flow.tofile(f)
+    f.close()
+
+
+def segment_flow(flow):
+    h = flow.shape[0]
+    w = flow.shape[1]
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW)
+    idx2 = abs(u) == SMALLFLOW
+    class0 = (v == 0) & (u == 0)
+    u[idx2] = 0.00001
+    tan_value = v / u
+
+    class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0)
+    class2 = (tan_value >= 1) & (u >= 0) & (v >= 0)
+    class3 = (tan_value < -1) & (u <= 0) & (v >= 0)
+    class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0)
+    class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0)
+    class7 = (tan_value < -1) & (u >= 0) & (v <= 0)
+    class6 = (tan_value >= 1) & (u <= 0) & (v <= 0)
+    class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0)
+
+    seg = np.zeros((h, w))
+
+    seg[class1] = 1
+    seg[class2] = 2
+    seg[class3] = 3
+    seg[class4] = 4
+    seg[class5] = 5
+    seg[class6] = 6
+    seg[class7] = 7
+    seg[class8] = 8
+    seg[class0] = 0
+    seg[idx] = 0
+
+    return seg
+
+
+def flow_error(tu, tv, u, v):
+    """
+    Calculate average end point error
+    :param tu: ground-truth horizontal flow map
+    :param tv: ground-truth vertical flow map
+    :param u:  estimated horizontal flow map
+    :param v:  estimated vertical flow map
+    :return: End point error of the estimated flow
+    """
+    smallflow = 0.0
+    """
+    stu = tu[bord+1:end-bord,bord+1:end-bord]
+    stv = tv[bord+1:end-bord,bord+1:end-bord]
+    su = u[bord+1:end-bord,bord+1:end-bord]
+    sv = v[bord+1:end-bord,bord+1:end-bord]
+    """
+    stu = tu[:]
+    stv = tv[:]
+    su = u[:]
+    sv = v[:]
+
+    idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH)
+    stu[idxUnknow] = 0
+    stv[idxUnknow] = 0
+    su[idxUnknow] = 0
+    sv[idxUnknow] = 0
+
+    ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)]
+    index_su = su[ind2]
+    index_sv = sv[ind2]
+    an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1)
+    un = index_su * an
+    vn = index_sv * an
+
+    index_stu = stu[ind2]
+    index_stv = stv[ind2]
+    tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1)
+    tun = index_stu * tn
+    tvn = index_stv * tn
+
+    """
+    angle = un * tun + vn * tvn + (an * tn)
+    index = [angle == 1.0]
+    angle[index] = 0.999
+    ang = np.arccos(angle)
+    mang = np.mean(ang)
+    mang = mang * 180 / np.pi
+    """
+
+    epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2)
+    epe = epe[ind2]
+    mepe = np.mean(epe)
+    return mepe
+
+
+def flow_to_image(flow, display=False):
+    """
+    Convert flow into middlebury color code image
+    :param flow: optical flow map
+    :return: optical flow image in middlebury color
+    """
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    maxu = -999.0
+    maxv = -999.0
+    minu = 999.0
+    minv = 999.0
+
+    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
+    u[idxUnknow] = 0
+    v[idxUnknow] = 0
+
+    maxu = max(maxu, np.max(u))
+    minu = min(minu, np.min(u))
+
+    maxv = max(maxv, np.max(v))
+    minv = min(minv, np.min(v))
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+    maxrad = max(-1, np.max(rad))
+
+    if display:
+        print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv))
+
+    u = u / (maxrad + np.finfo(float).eps)
+    v = v / (maxrad + np.finfo(float).eps)
+
+    img = compute_color(u, v)
+
+    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
+    img[idx] = 0
+
+    return np.uint8(img)
+
+
+def evaluate_flow_file(gt, pred):
+    """
+    evaluate the estimated optical flow end point error according to ground truth provided
+    :param gt: ground truth file path
+    :param pred: estimated optical flow file path
+    :return: end point error, float32
+    """
+    # Read flow files and calculate the errors
+    gt_flow = read_flow(gt)  # ground truth flow
+    eva_flow = read_flow(pred)  # predicted flow
+    # Calculate errors
+    average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1])
+    return average_pe
+
+
+def evaluate_flow(gt_flow, pred_flow):
+    """
+    gt: ground-truth flow
+    pred: estimated flow
+    """
+    average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], pred_flow[:, :, 0], pred_flow[:, :, 1])
+    return average_pe
+
+
+"""
+==============
+Disparity Section
+==============
+"""
+
+
+def read_disp_png(file_name):
+    """
+    Read optical flow from KITTI .png file
+    :param file_name: name of the flow file
+    :return: optical flow data in matrix
+    """
+    image_object = png.Reader(filename=file_name)
+    image_direct = image_object.asDirect()
+    image_data = list(image_direct[2])
+    (w, h) = image_direct[3]["size"]
+    channel = len(image_data[0]) / w
+    flow = np.zeros((h, w, channel), dtype=np.uint16)
+    for i in range(len(image_data)):
+        for j in range(channel):
+            flow[i, :, j] = image_data[i][j::channel]
+    return flow[:, :, 0] / 256
+
+
+def disp_to_flowfile(disp, filename):
+    """
+    Read KITTI disparity file in png format
+    :param disp: disparity matrix
+    :param filename: the flow file name to save
+    :return: None
+    """
+    f = open(filename, "wb")
+    magic = np.array([202021.25], dtype=np.float32)
+    (height, width) = disp.shape[0:2]
+    w = np.array([width], dtype=np.int32)
+    h = np.array([height], dtype=np.int32)
+    empty_map = np.zeros((height, width), dtype=np.float32)
+    data = np.dstack((disp, empty_map))
+    magic.tofile(f)
+    w.tofile(f)
+    h.tofile(f)
+    data.tofile(f)
+    f.close()
+
+
+"""
+==============
+Image Section
+==============
+"""
+
+
+def read_image(filename):
+    """
+    Read normal image of any format
+    :param filename: name of the image file
+    :return: image data in matrix uint8 type
+    """
+    img = Image.open(filename)
+    im = np.array(img)
+    return im
+
+
+def warp_image(im, flow):
+    """
+    Use optical flow to warp image to the next
+    :param im: image to warp
+    :param flow: optical flow
+    :return: warped image
+    """
+    from scipy import interpolate
+
+    image_height = im.shape[0]
+    image_width = im.shape[1]
+    flow_height = flow.shape[0]
+    flow_width = flow.shape[1]
+    n = image_height * image_width
+    (iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width])
+    (fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width])
+    fx += flow[:, :, 0].astype(np.float)
+    fy += flow[:, :, 1]
+    mask = np.logical_or(fx < 0, fx > flow_width)
+    mask = np.logical_or(mask, fy < 0)
+    mask = np.logical_or(mask, fy > flow_height)
+    fx = np.minimum(np.maximum(fx, 0), flow_width)
+    fy = np.minimum(np.maximum(fy, 0), flow_height)
+    points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1)
+    xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1)
+    warp = np.zeros((image_height, image_width, im.shape[2]))
+    for i in range(im.shape[2]):
+        channel = im[:, :, i]
+        # plt.imshow(channel, cmap='gray')
+        values = channel.reshape(n, 1)
+        new_channel = interpolate.griddata(points, values, xi, method="cubic")
+        new_channel = np.reshape(new_channel, [flow_height, flow_width])
+        new_channel[mask] = 1
+        warp[:, :, i] = new_channel.astype(np.uint8)
+
+    return warp.astype(np.uint8)
+
+
+"""
+==============
+Others
+==============
+"""
+
+
+def scale_image(image, new_range):
+    """
+    Linearly scale the image into desired range
+    :param image: input image
+    :param new_range: the new range to be aligned
+    :return: image normalized in new range
+    """
+    min_val = np.min(image).astype(np.float32)
+    max_val = np.max(image).astype(np.float32)
+    min_val_new = np.array(min(new_range), dtype=np.float32)
+    max_val_new = np.array(max(new_range), dtype=np.float32)
+    scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new
+    return scaled_image.astype(np.uint8)
+
+
+def compute_color(u, v):
+    """
+    compute optical flow color map
+    :param u: optical flow horizontal map
+    :param v: optical flow vertical map
+    :return: optical flow in color code
+    """
+    [h, w] = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+
+    colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+
+    a = np.arctan2(-v, -u) / np.pi
+
+    fk = (a + 1) / 2 * (ncols - 1) + 1
+
+    k0 = np.floor(fk).astype(int)
+
+    k1 = k0 + 1
+    k1[k1 == ncols + 1] = 1
+    f = fk - k0
+
+    for i in range(0, np.size(colorwheel, 1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0 - 1] / 255
+        col1 = tmp[k1 - 1] / 255
+        col = (1 - f) * col0 + f * col1
+
+        idx = rad <= 1
+        col[idx] = 1 - rad[idx] * (1 - col[idx])
+        notidx = np.logical_not(idx)
+
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))
+
+    return img
+
+
+def make_color_wheel():
+    """
+    Generate color wheel according Middlebury color code
+    :return: Color wheel
+    """
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+
+    colorwheel = np.zeros([ncols, 3])
+
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
+    col += RY
+
+    # YG
+    colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG))
+    colorwheel[col : col + YG, 1] = 255
+    col += YG
+
+    # GC
+    colorwheel[col : col + GC, 1] = 255
+    colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))
+    col += GC
+
+    # CB
+    colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB))
+    colorwheel[col : col + CB, 2] = 255
+    col += CB
+
+    # BM
+    colorwheel[col : col + BM, 2] = 255
+    colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))
+    col += +BM
+
+    # MR
+    colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col : col + MR, 0] = 255
+
+    return colorwheel

lib/functional.py

@@ -0,0 +1,606 @@
+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
+
+
+def _is_pil_image(img):
+    if accimage is not None:
+        return isinstance(img, (Image.Image, accimage.Image))
+    else:
+        return isinstance(img, Image.Image)
+
+
+def _is_tensor_image(img):
+    return torch.is_tensor(img) and img.ndimension() == 3
+
+
+def _is_numpy_image(img):
+    return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
+
+
+def to_tensor(pic):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
+
+    See ``ToTensor`` for more details.
+
+    Args:
+        pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+    Returns:
+        Tensor: Converted image.
+    """
+    if not (_is_pil_image(pic) or _is_numpy_image(pic)):
+        raise TypeError("pic should be PIL Image or ndarray. Got {}".format(type(pic)))
+
+    if isinstance(pic, np.ndarray):
+        # handle numpy array
+        img = torch.from_numpy(pic.transpose((2, 0, 1)))
+        # backward compatibility
+        return img.float().div(255)
+
+    if accimage is not None and isinstance(pic, accimage.Image):
+        nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
+        pic.copyto(nppic)
+        return torch.from_numpy(nppic)
+
+    # handle PIL Image
+    if pic.mode == "I":
+        img = torch.from_numpy(np.array(pic, np.int32, copy=False))
+    elif pic.mode == "I;16":
+        img = torch.from_numpy(np.array(pic, np.int16, copy=False))
+    else:
+        img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
+    # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
+    if pic.mode == "YCbCr":
+        nchannel = 3
+    elif pic.mode == "I;16":
+        nchannel = 1
+    else:
+        nchannel = len(pic.mode)
+    img = img.view(pic.size[1], pic.size[0], nchannel)
+    # put it from HWC to CHW format
+    # yikes, this transpose takes 80% of the loading time/CPU
+    img = img.transpose(0, 1).transpose(0, 2).contiguous()
+    if isinstance(img, torch.ByteTensor):
+        return img.float().div(255)
+    else:
+        return img
+
+
+def to_mytensor(pic):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
+
+    See ``ToTensor`` for more details.
+
+    Args:
+        pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+    Returns:
+        Tensor: Converted image.
+    """
+    pic_arr = np.array(pic)
+    if pic_arr.ndim == 2:
+        pic_arr = pic_arr[..., np.newaxis]
+    img = torch.from_numpy(pic_arr.transpose((2, 0, 1)))
+    if not isinstance(img, torch.FloatTensor):
+        return img.float()  # no normalize .div(255)
+    else:
+        return img
+
+
+def to_pil_image(pic, mode=None):
+    """Convert a tensor or an ndarray to PIL Image.
+
+    See :class:`~torchvision.transforms.ToPIlImage` for more details.
+
+    Args:
+        pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.
+        mode (`PIL.Image mode`_): color space and pixel depth of input data (optional).
+
+    .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes
+
+    Returns:
+        PIL Image: Image converted to PIL Image.
+    """
+    if not (_is_numpy_image(pic) or _is_tensor_image(pic)):
+        raise TypeError("pic should be Tensor or ndarray. Got {}.".format(type(pic)))
+
+    npimg = pic
+    if isinstance(pic, torch.FloatTensor):
+        pic = pic.mul(255).byte()
+    if torch.is_tensor(pic):
+        npimg = np.transpose(pic.numpy(), (1, 2, 0))
+
+    if not isinstance(npimg, np.ndarray):
+        raise TypeError("Input pic must be a torch.Tensor or NumPy ndarray, " + "not {}".format(type(npimg)))
+
+    if npimg.shape[2] == 1:
+        expected_mode = None
+        npimg = npimg[:, :, 0]
+        if npimg.dtype == np.uint8:
+            expected_mode = "L"
+        if npimg.dtype == np.int16:
+            expected_mode = "I;16"
+        if npimg.dtype == np.int32:
+            expected_mode = "I"
+        elif npimg.dtype == np.float32:
+            expected_mode = "F"
+        if mode is not None and mode != expected_mode:
+            raise ValueError(
+                "Incorrect mode ({}) supplied for input type {}. Should be {}".format(mode, np.dtype, expected_mode)
+            )
+        mode = expected_mode
+
+    elif npimg.shape[2] == 4:
+        permitted_4_channel_modes = ["RGBA", "CMYK"]
+        if mode is not None and mode not in permitted_4_channel_modes:
+            raise ValueError("Only modes {} are supported for 4D inputs".format(permitted_4_channel_modes))
+
+        if mode is None and npimg.dtype == np.uint8:
+            mode = "RGBA"
+    else:
+        permitted_3_channel_modes = ["RGB", "YCbCr", "HSV"]
+        if mode is not None and mode not in permitted_3_channel_modes:
+            raise ValueError("Only modes {} are supported for 3D inputs".format(permitted_3_channel_modes))
+        if mode is None and npimg.dtype == np.uint8:
+            mode = "RGB"
+
+    if mode is None:
+        raise TypeError("Input type {} is not supported".format(npimg.dtype))
+
+    return Image.fromarray(npimg, mode=mode)
+
+
+def normalize(tensor, mean, std):
+    """Normalize a tensor image with mean and standard deviation.
+
+    See ``Normalize`` for more details.
+
+    Args:
+        tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channely.
+
+    Returns:
+        Tensor: Normalized Tensor image.
+    """
+    if not _is_tensor_image(tensor):
+        raise TypeError("tensor is not a torch image.")
+    # TODO: make efficient
+    if tensor.size(0) == 1:
+        tensor.sub_(mean).div_(std)
+    else:
+        for t, m, s in zip(tensor, mean, std):
+            t.sub_(m).div_(s)
+    return tensor
+
+
+def resize(img, size, interpolation=Image.BILINEAR):
+    """Resize the input PIL Image to the given size.
+
+    Args:
+        img (PIL Image): Image to be resized.
+        size (sequence or int): Desired output size. If size is a sequence like
+            (h, w), the output size will be matched to this. If size is an int,
+            the smaller edge of the image will be matched to this number maintaing
+            the aspect ratio. 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``
+
+    Returns:
+        PIL Image: Resized image.
+    """
+    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)):
+        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:
+        return img.resize(size[::-1], interpolation)
+
+
+def scale(*args, **kwargs):
+    warnings.warn("The use of the transforms.Scale transform is deprecated, " + "please use transforms.Resize instead.")
+    return resize(*args, **kwargs)
+
+
+def pad(img, padding, fill=0):
+    """Pad the given PIL Image on all sides with the given "pad" value.
+
+    Args:
+        img (PIL Image): Image to be padded.
+        padding (int or tuple): Padding on each border. If a single int is provided this
+            is used to pad all borders. If tuple of length 2 is provided this is the padding
+            on left/right and top/bottom respectively. If a tuple of length 4 is provided
+            this is the padding for the left, top, right and bottom borders
+            respectively.
+        fill: Pixel fill value. Default is 0. If a tuple of
+            length 3, it is used to fill R, G, B channels respectively.
+
+    Returns:
+        PIL Image: Padded image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    if not isinstance(padding, (numbers.Number, tuple)):
+        raise TypeError("Got inappropriate padding arg")
+    if not isinstance(fill, (numbers.Number, str, tuple)):
+        raise TypeError("Got inappropriate fill arg")
+
+    if isinstance(padding, collections.Sequence) and len(padding) not in [2, 4]:
+        raise ValueError(
+            "Padding must be an int or a 2, or 4 element tuple, not a " + "{} element tuple".format(len(padding))
+        )
+
+    return ImageOps.expand(img, border=padding, fill=fill)
+
+
+def crop(img, i, j, h, w):
+    """Crop the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be cropped.
+        i: Upper pixel coordinate.
+        j: Left pixel coordinate.
+        h: Height of the cropped image.
+        w: Width of the cropped image.
+
+    Returns:
+        PIL Image: Cropped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    return img.crop((j, i, j + w, i + h))
+
+
+def center_crop(img, output_size):
+    if isinstance(output_size, numbers.Number):
+        output_size = (int(output_size), int(output_size))
+    w, h = img.size
+    th, tw = output_size
+    i = int(round((h - th) / 2.0))
+    j = int(round((w - tw) / 2.0))
+    return crop(img, i, j, th, tw)
+
+
+def resized_crop(img, i, j, h, w, size, interpolation=Image.BILINEAR):
+    """Crop the given PIL Image and resize it to desired size.
+
+    Notably used in RandomResizedCrop.
+
+    Args:
+        img (PIL Image): Image to be cropped.
+        i: Upper pixel coordinate.
+        j: Left pixel coordinate.
+        h: Height of the cropped image.
+        w: Width of the cropped image.
+        size (sequence or int): Desired output size. Same semantics as ``scale``.
+        interpolation (int, optional): Desired interpolation. Default is
+            ``PIL.Image.BILINEAR``.
+    Returns:
+        PIL Image: Cropped image.
+    """
+    assert _is_pil_image(img), "img should be PIL Image"
+    img = crop(img, i, j, h, w)
+    img = resize(img, size, interpolation)
+    return img
+
+
+def hflip(img):
+    """Horizontally flip the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be flipped.
+
+    Returns:
+        PIL Image:  Horizontall flipped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    return img.transpose(Image.FLIP_LEFT_RIGHT)
+
+
+def vflip(img):
+    """Vertically flip the given PIL Image.
+
+    Args:
+        img (PIL Image): Image to be flipped.
+
+    Returns:
+        PIL Image:  Vertically flipped image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    return img.transpose(Image.FLIP_TOP_BOTTOM)
+
+
+def five_crop(img, size):
+    """Crop the given PIL Image into four corners and the central crop.
+
+    .. Note::
+        This transform returns a tuple of images and there may be a
+        mismatch in the number of inputs and targets your ``Dataset`` returns.
+
+    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.
+    Returns:
+        tuple: tuple (tl, tr, bl, br, center) corresponding top left,
+            top right, bottom left, bottom right and center crop.
+    """
+    if isinstance(size, numbers.Number):
+        size = (int(size), int(size))
+    else:
+        assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+
+    w, h = img.size
+    crop_h, crop_w = size
+    if crop_w > w or crop_h > h:
+        raise ValueError("Requested crop size {} is bigger than input size {}".format(size, (h, w)))
+    tl = img.crop((0, 0, crop_w, crop_h))
+    tr = img.crop((w - crop_w, 0, w, crop_h))
+    bl = img.crop((0, h - crop_h, crop_w, h))
+    br = img.crop((w - crop_w, h - crop_h, w, h))
+    center = center_crop(img, (crop_h, crop_w))
+    return (tl, tr, bl, br, center)
+
+
+def ten_crop(img, size, vertical_flip=False):
+    """Crop the given PIL Image into four corners and the central crop plus the
+       flipped version of these (horizontal flipping is used by default).
+
+    .. Note::
+        This transform returns a tuple of images and there may be a
+        mismatch in the number of inputs and targets your ``Dataset`` returns.
+
+       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.
+           vertical_flip (bool): Use vertical flipping instead of horizontal
+
+        Returns:
+            tuple: tuple (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip,
+                br_flip, center_flip) corresponding top left, top right,
+                bottom left, bottom right and center crop and same for the
+                flipped image.
+    """
+    if isinstance(size, numbers.Number):
+        size = (int(size), int(size))
+    else:
+        assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+
+    first_five = five_crop(img, size)
+
+    if vertical_flip:
+        img = vflip(img)
+    else:
+        img = hflip(img)
+
+    second_five = five_crop(img, size)
+    return first_five + second_five
+
+
+def adjust_brightness(img, brightness_factor):
+    """Adjust brightness of an Image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        brightness_factor (float):  How much to adjust the brightness. Can be
+            any non negative number. 0 gives a black image, 1 gives the
+            original image while 2 increases the brightness by a factor of 2.
+
+    Returns:
+        PIL Image: Brightness adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    enhancer = ImageEnhance.Brightness(img)
+    img = enhancer.enhance(brightness_factor)
+    return img
+
+
+def adjust_contrast(img, contrast_factor):
+    """Adjust contrast of an Image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        contrast_factor (float): How much to adjust the contrast. Can be any
+            non negative number. 0 gives a solid gray image, 1 gives the
+            original image while 2 increases the contrast by a factor of 2.
+
+    Returns:
+        PIL Image: Contrast adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    enhancer = ImageEnhance.Contrast(img)
+    img = enhancer.enhance(contrast_factor)
+    return img
+
+
+def adjust_saturation(img, saturation_factor):
+    """Adjust color saturation of an image.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        saturation_factor (float):  How much to adjust the saturation. 0 will
+            give a black and white image, 1 will give the original image while
+            2 will enhance the saturation by a factor of 2.
+
+    Returns:
+        PIL Image: Saturation adjusted image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    enhancer = ImageEnhance.Color(img)
+    img = enhancer.enhance(saturation_factor)
+    return img
+
+
+def adjust_hue(img, hue_factor):
+    """Adjust hue of an image.
+
+    The image hue is adjusted by converting the image to HSV and
+    cyclically shifting the intensities in the hue channel (H).
+    The image is then converted back to original image mode.
+
+    `hue_factor` is the amount of shift in H channel and must be in the
+    interval `[-0.5, 0.5]`.
+
+    See https://en.wikipedia.org/wiki/Hue for more details on Hue.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        hue_factor (float):  How much to shift the hue channel. Should be in
+            [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in
+            HSV space in positive and negative direction respectively.
+            0 means no shift. Therefore, both -0.5 and 0.5 will give an image
+            with complementary colors while 0 gives the original image.
+
+    Returns:
+        PIL Image: Hue adjusted image.
+    """
+    if not (-0.5 <= hue_factor <= 0.5):
+        raise ValueError("hue_factor is not in [-0.5, 0.5].".format(hue_factor))
+
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    input_mode = img.mode
+    if input_mode in {"L", "1", "I", "F"}:
+        return img
+
+    h, s, v = img.convert("HSV").split()
+
+    np_h = np.array(h, dtype=np.uint8)
+    # uint8 addition take cares of rotation across boundaries
+    with np.errstate(over="ignore"):
+        np_h += np.uint8(hue_factor * 255)
+    h = Image.fromarray(np_h, "L")
+
+    img = Image.merge("HSV", (h, s, v)).convert(input_mode)
+    return img
+
+
+def adjust_gamma(img, gamma, gain=1):
+    """Perform gamma correction on an image.
+
+    Also known as Power Law Transform. Intensities in RGB mode are adjusted
+    based on the following equation:
+
+        I_out = 255 * gain * ((I_in / 255) ** gamma)
+
+    See https://en.wikipedia.org/wiki/Gamma_correction for more details.
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        gamma (float): Non negative real number. gamma larger than 1 make the
+            shadows darker, while gamma smaller than 1 make dark regions
+            lighter.
+        gain (float): The constant multiplier.
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    if gamma < 0:
+        raise ValueError("Gamma should be a non-negative real number")
+
+    input_mode = img.mode
+    img = img.convert("RGB")
+
+    np_img = np.array(img, dtype=np.float32)
+    np_img = 255 * gain * ((np_img / 255) ** gamma)
+    np_img = np.uint8(np.clip(np_img, 0, 255))
+
+    img = Image.fromarray(np_img, "RGB").convert(input_mode)
+    return img
+
+
+def rotate(img, angle, resample=False, expand=False, center=None):
+    """Rotate the image by angle and then (optionally) translate it by (n_columns, n_rows)
+
+
+    Args:
+        img (PIL Image): PIL Image to be rotated.
+        angle ({float, int}): In degrees degrees counter clockwise order.
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        expand (bool, optional): Optional expansion flag.
+            If true, expands the output image to make it large enough to hold the entire rotated image.
+            If false or omitted, make the output image the same size as the input image.
+            Note that the expand flag assumes rotation around the center and no translation.
+        center (2-tuple, optional): Optional center of rotation.
+            Origin is the upper left corner.
+            Default is the center of the image.
+    """
+
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    return img.rotate(angle, resample, expand, center)
+
+
+def to_grayscale(img, num_output_channels=1):
+    """Convert image to grayscale version of image.
+
+    Args:
+        img (PIL Image): Image to be converted to grayscale.
+
+    Returns:
+        PIL Image:  Grayscale version of the image.
+                    if num_output_channels == 1 : returned image is single channel
+                    if num_output_channels == 3 : returned image is 3 channel with r == g == b
+    """
+    if not _is_pil_image(img):
+        raise TypeError("img should be PIL Image. Got {}".format(type(img)))
+
+    if num_output_channels == 1:
+        img = img.convert("L")
+    elif num_output_channels == 3:
+        img = img.convert("L")
+        np_img = np.array(img, dtype=np.uint8)
+        np_img = np.dstack([np_img, np_img, np_img])
+        img = Image.fromarray(np_img, "RGB")
+    else:
+        raise ValueError("num_output_channels should be either 1 or 3")
+
+    return img

lib/imageloader_imagenet.py

@@ -0,0 +1,179 @@
+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

lib/imports.py

@@ -0,0 +1,82 @@
+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
+

lib/videoloader.py

@@ -0,0 +1,390 @@
+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

lib/videoloader_imagenet.py

@@ -0,0 +1,203 @@
+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