refactor
This commit is contained in:
zhangmozhe
@@ -1,390 +0,0 @@
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, "..")
|
||||
import os
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
import pdb
|
||||
import random
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import torch
|
||||
import torchvision
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from skimage import color
|
||||
from utils.flowlib import read_flow
|
||||
from utils.util_distortion import CenterPad, CenterPadCrop_numpy
|
||||
from utils.warping import WarpingLayer
|
||||
|
||||
import lib.functional as F
|
||||
|
||||
|
||||
class RGB2Lab(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
image_lab = color.rgb2lab(inputs)
|
||||
return image_lab
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
|
||||
inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
|
||||
return inputs
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inputs):
|
||||
outputs = F.to_mytensor(inputs) # permute channel and transform to tensor
|
||||
return outputs
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
def __init__(self, probability=0.6, sl=0.05, sh=0.6):
|
||||
self.probability = probability
|
||||
self.sl = sl
|
||||
self.sh = sh
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
if random.uniform(0, 1) > self.probability:
|
||||
return Image.fromarray(img)
|
||||
|
||||
area = img.shape[0] * img.shape[1]
|
||||
h0 = img.shape[0]
|
||||
w0 = img.shape[1]
|
||||
channel = img.shape[2]
|
||||
|
||||
h = int(round(random.uniform(self.sl, self.sh) * h0))
|
||||
w = int(round(random.uniform(self.sl, self.sh) * w0))
|
||||
|
||||
if w < img.shape[1] and h < img.shape[0]:
|
||||
x1 = random.randint(0, img.shape[0] - h)
|
||||
y1 = random.randint(0, img.shape[1] - w)
|
||||
img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
|
||||
return Image.fromarray(img)
|
||||
|
||||
return Image.fromarray(img)
|
||||
|
||||
|
||||
class CenterCrop(object):
|
||||
"""
|
||||
center crop the numpy array
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
# target size
|
||||
self.h0, self.w0 = image_size
|
||||
|
||||
def __call__(self, input_numpy):
|
||||
if input_numpy.ndim == 3:
|
||||
h, w, channel = input_numpy.shape
|
||||
output_numpy = np.zeros((self.h0, self.w0, channel))
|
||||
output_numpy = input_numpy[
|
||||
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, :
|
||||
]
|
||||
else:
|
||||
h, w = input_numpy.shape
|
||||
output_numpy = np.zeros((self.h0, self.w0))
|
||||
output_numpy = input_numpy[
|
||||
(h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0
|
||||
]
|
||||
return output_numpy
|
||||
|
||||
|
||||
def parse_images(data_root):
|
||||
image_pairs = []
|
||||
subdirs = sorted(os.listdir(data_root))
|
||||
for subdir in subdirs:
|
||||
path = os.path.join(data_root, subdir)
|
||||
if not os.path.isdir(path):
|
||||
continue
|
||||
|
||||
# parse_file = os.path.join(path, 'pairs_output.txt')
|
||||
parse_file = os.path.join(path, "pairs_output_new.txt")
|
||||
if os.path.exists(parse_file):
|
||||
f = open(parse_file, "r")
|
||||
lines = f.readlines()
|
||||
for line in lines:
|
||||
line = line.replace("\n", "")
|
||||
(
|
||||
image1_name,
|
||||
image2_name,
|
||||
reference_video_name,
|
||||
reference_video_name1,
|
||||
reference_name1,
|
||||
reference_name2,
|
||||
reference_name3,
|
||||
reference_name4,
|
||||
reference_name5,
|
||||
reference_gt1,
|
||||
reference_gt2,
|
||||
reference_gt3,
|
||||
) = line.split()
|
||||
image1_name = image1_name.split(".")[0]
|
||||
image2_name = image2_name.split(".")[0]
|
||||
reference_video_name = reference_video_name.split(".")[0]
|
||||
reference_video_name1 = reference_video_name1.split(".")[0]
|
||||
reference_name1 = reference_name1.split(".")[0]
|
||||
reference_name2 = reference_name2.split(".")[0]
|
||||
reference_name3 = reference_name3.split(".")[0]
|
||||
reference_name4 = reference_name4.split(".")[0]
|
||||
reference_name5 = reference_name5.split(".")[0]
|
||||
|
||||
reference_gt1 = reference_gt1.split(".")[0]
|
||||
reference_gt2 = reference_gt2.split(".")[0]
|
||||
reference_gt3 = reference_gt3.split(".")[0]
|
||||
|
||||
flow_forward_name = image1_name + "_forward"
|
||||
flow_backward_name = image1_name + "_backward"
|
||||
mask_name = image1_name + "_mask"
|
||||
|
||||
item = (
|
||||
image1_name + ".jpg",
|
||||
image2_name + ".jpg",
|
||||
reference_video_name + ".jpg",
|
||||
reference_name1 + ".JPEG",
|
||||
reference_name2 + ".JPEG",
|
||||
reference_name3 + ".JPEG",
|
||||
reference_name4 + ".JPEG",
|
||||
reference_name5 + ".JPEG",
|
||||
flow_forward_name + ".flo",
|
||||
flow_backward_name + ".flo",
|
||||
mask_name + ".pgm",
|
||||
reference_gt1 + ".jpg",
|
||||
reference_gt2 + ".jpg",
|
||||
reference_gt3 + ".jpg",
|
||||
path,
|
||||
)
|
||||
image_pairs.append(item)
|
||||
|
||||
f.close()
|
||||
|
||||
else:
|
||||
raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n"))
|
||||
|
||||
return image_pairs
|
||||
|
||||
|
||||
class VideosDataset(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
epoch,
|
||||
image_size,
|
||||
image_transform=None,
|
||||
use_google_reference=False,
|
||||
real_reference_probability=1,
|
||||
nonzero_placeholder_probability=0.5,
|
||||
):
|
||||
self.data_root = data_root
|
||||
self.image_transform = image_transform
|
||||
self.CenterPad = CenterPad(image_size)
|
||||
self.ToTensor = ToTensor()
|
||||
self.CenterCrop = CenterCrop(image_size)
|
||||
|
||||
assert len(self.data_root) > 0, "find no dataroot"
|
||||
self.epoch = epoch
|
||||
self.image_pairs = parse_images(self.data_root)
|
||||
self.real_len = len(self.image_pairs)
|
||||
print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len))
|
||||
self.image_pairs *= epoch
|
||||
self.use_google_reference = use_google_reference
|
||||
self.real_reference_probability = real_reference_probability
|
||||
self.nonzero_placeholder_probability = nonzero_placeholder_probability
|
||||
|
||||
def __getitem__(self, index):
|
||||
(
|
||||
image1_name,
|
||||
image2_name,
|
||||
reference_video_name,
|
||||
reference_name1,
|
||||
reference_name2,
|
||||
reference_name3,
|
||||
reference_name4,
|
||||
reference_name5,
|
||||
flow_forward_name,
|
||||
flow_backward_name,
|
||||
mask_name,
|
||||
reference_gt1,
|
||||
reference_gt2,
|
||||
reference_gt3,
|
||||
path,
|
||||
) = self.image_pairs[index]
|
||||
try:
|
||||
# print('video loader')
|
||||
I1 = Image.open(os.path.join(path, "input_pad", image1_name))
|
||||
I2 = Image.open(os.path.join(path, "input_pad", image2_name))
|
||||
|
||||
# randomly choose a reference frame from the video
|
||||
I_reference_video = Image.open(
|
||||
os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3]))
|
||||
)
|
||||
I_reference_video_real = Image.open(
|
||||
os.path.join(
|
||||
path,
|
||||
"reference",
|
||||
random.choice(
|
||||
[reference_name1, reference_name2, reference_name3, reference_name4, reference_name5]
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name)) # numpy
|
||||
flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name)) # numpy
|
||||
mask = Image.open(os.path.join(path, "mask", mask_name))
|
||||
|
||||
# binary mask
|
||||
mask = np.array(mask)
|
||||
mask[mask < 240] = 0
|
||||
mask[mask >= 240] = 1
|
||||
|
||||
# transform
|
||||
I1 = self.image_transform(I1)
|
||||
I2 = self.image_transform(I2)
|
||||
I_reference_video = self.image_transform(self.CenterPad(I_reference_video))
|
||||
I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real))
|
||||
flow_forward = self.ToTensor(self.CenterCrop(flow_forward))
|
||||
flow_backward = self.ToTensor(self.CenterCrop(flow_backward))
|
||||
mask = self.ToTensor(self.CenterCrop(mask))
|
||||
|
||||
if np.random.random() < self.real_reference_probability:
|
||||
I_reference_output = I_reference_video_real
|
||||
placeholder = torch.zeros_like(I1)
|
||||
self_ref_flag = torch.zeros_like(I1)
|
||||
else:
|
||||
I_reference_output = I_reference_video
|
||||
placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
|
||||
self_ref_flag = torch.ones_like(I1)
|
||||
|
||||
outputs = []
|
||||
outputs.append(I1)
|
||||
outputs.append(I2)
|
||||
outputs.append(I_reference_output)
|
||||
# outputs.append(I_reference_video_real)
|
||||
outputs.append(flow_forward)
|
||||
outputs.append(flow_backward)
|
||||
outputs.append(mask)
|
||||
outputs.append(placeholder)
|
||||
outputs.append(self_ref_flag)
|
||||
|
||||
except Exception as e:
|
||||
print("problem in, ", path)
|
||||
print(e)
|
||||
return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
|
||||
|
||||
return outputs
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_pairs)
|
||||
|
||||
|
||||
import torchvision.utils as vutils
|
||||
from torch.autograd import Variable
|
||||
|
||||
|
||||
def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8):
|
||||
if isinstance(img_l_mc, Variable):
|
||||
img_l_mc = img_l_mc.data.cpu()
|
||||
if isinstance(img_ab_mc, Variable):
|
||||
img_ab_mc = img_ab_mc.data.cpu()
|
||||
|
||||
if img_l_mc.is_cuda:
|
||||
img_l_mc = img_l_mc.cpu()
|
||||
if img_ab_mc.is_cuda:
|
||||
img_ab_mc = img_ab_mc.cpu()
|
||||
|
||||
assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input"
|
||||
|
||||
l_norm, ab_norm = 1.0, 1.0
|
||||
l_mean, ab_mean = 50.0, 0
|
||||
img_l = img_l_mc * l_norm + l_mean
|
||||
img_ab = img_ab_mc * ab_norm + ab_mean
|
||||
pred_lab = torch.cat((img_l, img_ab), dim=1)
|
||||
grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64")
|
||||
grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")
|
||||
return grid_rgb
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
def test_videoloader():
|
||||
# dataset initI1lization
|
||||
# image size = 256
|
||||
train_dataset = VideosDataset(data_root="../../Dataset/video/video_pair", epoch=10)
|
||||
|
||||
data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
|
||||
|
||||
iter_instance = iter(data_loader)
|
||||
sample = iter_instance.next()
|
||||
I1_lab, I2_lab, I_reference_lab, flow_forward, flow_backward, mask = sample
|
||||
|
||||
# l: [-50,50], ab: [-128, 128]
|
||||
l_norm, ab_norm = 1.0, 1.0
|
||||
l_mean, ab_mean = 50.0, 0
|
||||
I1_l = I1_lab[:, 0:1, :, :]
|
||||
I1_ab = I1_lab[:, 1:3, :, :]
|
||||
I2_l = I2_lab[:, 0:1, :, :]
|
||||
I2_ab = I2_lab[:, 1:3, :, :]
|
||||
I_reference_l = I_reference_lab[:, 0:1, :, :]
|
||||
I_reference_ab = I_reference_lab[:, 1:3, :, :]
|
||||
|
||||
# test image samples
|
||||
image_A_grid = batch_lab2rgb_transpose_mc(I1_l, I1_ab, nrow=4).astype(np.uint8)
|
||||
image_B_grid = batch_lab2rgb_transpose_mc(I2_l, I2_ab, nrow=4).astype(np.uint8)
|
||||
image_ref_grid = batch_lab2rgb_transpose_mc(I_reference_l, I_reference_ab, nrow=4).astype(np.uint8)
|
||||
plt.imsave("A.png", image_A_grid)
|
||||
plt.imsave("B.png", image_B_grid)
|
||||
plt.figure()
|
||||
plt.imshow(image_A_grid)
|
||||
plt.title("image A")
|
||||
plt.figure()
|
||||
plt.imshow(image_B_grid)
|
||||
plt.title("image B")
|
||||
plt.figure()
|
||||
plt.imshow(image_ref_grid)
|
||||
plt.title("image ref")
|
||||
# plt.show()
|
||||
|
||||
# test warping
|
||||
print(I2_lab.shape)
|
||||
print(flow_forward.shape)
|
||||
warping_layer = WarpingLayer("gpu")
|
||||
I21_lab = warping_layer(I2_lab.cuda(), flow_forward.cuda())
|
||||
I21_lab = I21_lab * mask.cuda()
|
||||
|
||||
I21_l = I21_lab[:, 0:1, :, :]
|
||||
I21_ab = I21_lab[:, 1:3, :, :]
|
||||
image_warp_grid = batch_lab2rgb_transpose_mc(I21_l, I21_ab, nrow=4).astype(np.uint8)
|
||||
plt.imsave("BA.png", image_warp_grid)
|
||||
plt.figure()
|
||||
plt.imshow(image_warp_grid)
|
||||
plt.title("image warp")
|
||||
# plt.show()
|
||||
|
||||
# image difference
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(image_warp_grid - image_A_grid).astype(np.uint8))
|
||||
plt.title("warp-A")
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(image_B_grid - image_A_grid).astype(np.uint8))
|
||||
plt.title("B-A")
|
||||
# plt.show()
|
||||
plt.imsave("difference_warp.png", np.abs(image_warp_grid - image_A_grid).astype(np.uint8))
|
||||
plt.imsave("difference_BA.png", np.abs(image_B_grid - image_A_grid).astype(np.uint8))
|
||||
|
||||
return None
|
||||
Reference in New Issue
Block a user