302 lines
10 KiB
Python
302 lines
10 KiB
Python
import sys
|
|
|
|
sys.path.insert(0, "..")
|
|
import os
|
|
import random
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.utils as vutils
|
|
from PIL import Image
|
|
from skimage import color
|
|
from torch.autograd import Variable
|
|
from utils.flowlib import read_flow
|
|
from utils.util_distortion import CenterPad
|
|
|
|
import lib.functional as F
|
|
|
|
cv2.setNumThreads(0)
|
|
|
|
|
|
class RGB2Lab(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
return color.rgb2lab(inputs)
|
|
|
|
|
|
class Normalize(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
|
|
inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
|
|
return inputs
|
|
|
|
|
|
class ToTensor(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
outputs = F.to_mytensor(inputs) # permute channel and transform to tensor
|
|
return outputs
|
|
|
|
|
|
class RandomErasing(object):
|
|
def __init__(self, probability=0.6, sl=0.05, sh=0.6):
|
|
self.probability = probability
|
|
self.sl = sl
|
|
self.sh = sh
|
|
|
|
def __call__(self, img):
|
|
img = np.array(img)
|
|
if random.uniform(0, 1) > self.probability:
|
|
return Image.fromarray(img)
|
|
|
|
area = img.shape[0] * img.shape[1]
|
|
h0 = img.shape[0]
|
|
w0 = img.shape[1]
|
|
channel = img.shape[2]
|
|
|
|
h = int(round(random.uniform(self.sl, self.sh) * h0))
|
|
w = int(round(random.uniform(self.sl, self.sh) * w0))
|
|
|
|
if w < img.shape[1] and h < img.shape[0]:
|
|
x1 = random.randint(0, img.shape[0] - h)
|
|
y1 = random.randint(0, img.shape[1] - w)
|
|
img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
|
|
return Image.fromarray(img)
|
|
|
|
return Image.fromarray(img)
|
|
|
|
|
|
class CenterCrop(object):
|
|
"""
|
|
center crop the numpy array
|
|
"""
|
|
|
|
def __init__(self, image_size):
|
|
self.h0, self.w0 = image_size
|
|
|
|
def __call__(self, input_numpy):
|
|
if input_numpy.ndim == 3:
|
|
h, w, channel = input_numpy.shape
|
|
output_numpy = np.zeros((self.h0, self.w0, channel))
|
|
output_numpy = input_numpy[
|
|
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, :
|
|
]
|
|
else:
|
|
h, w = input_numpy.shape
|
|
output_numpy = np.zeros((self.h0, self.w0))
|
|
output_numpy = input_numpy[
|
|
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0
|
|
]
|
|
return output_numpy
|
|
|
|
|
|
def parse_images(data_root):
|
|
image_pairs = []
|
|
subdirs = sorted(os.listdir(data_root))
|
|
for subdir in subdirs:
|
|
path = os.path.join(data_root, subdir)
|
|
if not os.path.isdir(path):
|
|
continue
|
|
|
|
parse_file = os.path.join(path, "pairs_output_new.txt")
|
|
if os.path.exists(parse_file):
|
|
with open(parse_file, "r") as f:
|
|
lines = f.readlines()
|
|
for line in lines:
|
|
line = line.replace("\n", "")
|
|
(
|
|
image1_name,
|
|
image2_name,
|
|
reference_video_name,
|
|
reference_video_name1,
|
|
reference_name1,
|
|
reference_name2,
|
|
reference_name3,
|
|
reference_name4,
|
|
reference_name5,
|
|
reference_gt1,
|
|
reference_gt2,
|
|
reference_gt3,
|
|
) = line.split()
|
|
image1_name = image1_name.split(".")[0]
|
|
image2_name = image2_name.split(".")[0]
|
|
reference_video_name = reference_video_name.split(".")[0]
|
|
reference_video_name1 = reference_video_name1.split(".")[0]
|
|
reference_name1 = reference_name1.split(".")[0]
|
|
reference_name2 = reference_name2.split(".")[0]
|
|
reference_name3 = reference_name3.split(".")[0]
|
|
reference_name4 = reference_name4.split(".")[0]
|
|
reference_name5 = reference_name5.split(".")[0]
|
|
|
|
reference_gt1 = reference_gt1.split(".")[0]
|
|
reference_gt2 = reference_gt2.split(".")[0]
|
|
reference_gt3 = reference_gt3.split(".")[0]
|
|
|
|
flow_forward_name = image1_name + "_forward"
|
|
flow_backward_name = image1_name + "_backward"
|
|
mask_name = image1_name + "_mask"
|
|
|
|
item = (
|
|
image1_name + ".jpg",
|
|
image2_name + ".jpg",
|
|
reference_video_name + ".jpg",
|
|
reference_name1 + ".JPEG",
|
|
reference_name2 + ".JPEG",
|
|
reference_name3 + ".JPEG",
|
|
reference_name4 + ".JPEG",
|
|
reference_name5 + ".JPEG",
|
|
flow_forward_name + ".flo",
|
|
flow_backward_name + ".flo",
|
|
mask_name + ".pgm",
|
|
reference_gt1 + ".jpg",
|
|
reference_gt2 + ".jpg",
|
|
reference_gt3 + ".jpg",
|
|
path,
|
|
)
|
|
image_pairs.append(item)
|
|
|
|
else:
|
|
raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n"))
|
|
|
|
return image_pairs
|
|
|
|
|
|
class VideosDataset(torch.utils.data.Dataset):
|
|
def __init__(
|
|
self,
|
|
data_root,
|
|
epoch,
|
|
image_size,
|
|
image_transform=None,
|
|
use_google_reference=False,
|
|
real_reference_probability=1,
|
|
nonzero_placeholder_probability=0.5,
|
|
):
|
|
self.data_root = data_root
|
|
self.image_transform = image_transform
|
|
self.CenterPad = CenterPad(image_size)
|
|
self.ToTensor = ToTensor()
|
|
self.CenterCrop = CenterCrop(image_size)
|
|
|
|
assert len(self.data_root) > 0, "find no dataroot"
|
|
self.epoch = epoch
|
|
self.image_pairs = parse_images(self.data_root)
|
|
self.real_len = len(self.image_pairs)
|
|
print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len))
|
|
self.image_pairs *= epoch
|
|
self.use_google_reference = use_google_reference
|
|
self.real_reference_probability = real_reference_probability
|
|
self.nonzero_placeholder_probability = nonzero_placeholder_probability
|
|
|
|
def __getitem__(self, index):
|
|
(
|
|
image1_name,
|
|
image2_name,
|
|
reference_video_name,
|
|
reference_name1,
|
|
reference_name2,
|
|
reference_name3,
|
|
reference_name4,
|
|
reference_name5,
|
|
flow_forward_name,
|
|
flow_backward_name,
|
|
mask_name,
|
|
reference_gt1,
|
|
reference_gt2,
|
|
reference_gt3,
|
|
path,
|
|
) = self.image_pairs[index]
|
|
try:
|
|
I1 = Image.open(os.path.join(path, "input_pad", image1_name))
|
|
I2 = Image.open(os.path.join(path, "input_pad", image2_name))
|
|
|
|
I_reference_video = Image.open(
|
|
os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3]))
|
|
)
|
|
I_reference_video_real = Image.open(
|
|
os.path.join(
|
|
path,
|
|
"reference",
|
|
random.choice(
|
|
[reference_name1, reference_name2, reference_name3, reference_name4, reference_name5]
|
|
),
|
|
)
|
|
)
|
|
|
|
flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name)) # numpy
|
|
flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name)) # numpy
|
|
mask = Image.open(os.path.join(path, "mask", mask_name))
|
|
|
|
# binary mask
|
|
mask = np.array(mask)
|
|
mask[mask < 240] = 0
|
|
mask[mask >= 240] = 1
|
|
|
|
# transform
|
|
I1 = self.image_transform(I1)
|
|
I2 = self.image_transform(I2)
|
|
I_reference_video = self.image_transform(self.CenterPad(I_reference_video))
|
|
I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real))
|
|
flow_forward = self.ToTensor(self.CenterCrop(flow_forward))
|
|
flow_backward = self.ToTensor(self.CenterCrop(flow_backward))
|
|
mask = self.ToTensor(self.CenterCrop(mask))
|
|
|
|
if np.random.random() < self.real_reference_probability:
|
|
I_reference_output = I_reference_video_real
|
|
placeholder = torch.zeros_like(I1)
|
|
self_ref_flag = torch.zeros_like(I1)
|
|
else:
|
|
I_reference_output = I_reference_video
|
|
placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
|
|
self_ref_flag = torch.ones_like(I1)
|
|
|
|
outputs = [
|
|
I1,
|
|
I2,
|
|
I_reference_output,
|
|
flow_forward,
|
|
flow_backward,
|
|
mask,
|
|
placeholder,
|
|
self_ref_flag,
|
|
]
|
|
|
|
except Exception as e:
|
|
print("problem in, ", path)
|
|
print(e)
|
|
return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
|
|
return outputs
|
|
|
|
def __len__(self):
|
|
return len(self.image_pairs)
|
|
|
|
|
|
def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8):
|
|
if isinstance(img_l_mc, Variable):
|
|
img_l_mc = img_l_mc.data.cpu()
|
|
if isinstance(img_ab_mc, Variable):
|
|
img_ab_mc = img_ab_mc.data.cpu()
|
|
|
|
if img_l_mc.is_cuda:
|
|
img_l_mc = img_l_mc.cpu()
|
|
if img_ab_mc.is_cuda:
|
|
img_ab_mc = img_ab_mc.cpu()
|
|
|
|
assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input"
|
|
|
|
l_norm, ab_norm = 1.0, 1.0
|
|
l_mean, ab_mean = 50.0, 0
|
|
img_l = img_l_mc * l_norm + l_mean
|
|
img_ab = img_ab_mc * ab_norm + ab_mean
|
|
pred_lab = torch.cat((img_l, img_ab), dim=1)
|
|
grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64")
|
|
return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|