|
import os.path |
|
import torch |
|
import torch.utils.data as data |
|
from PIL import Image |
|
import random |
|
import utils |
|
import numpy as np |
|
import torchvision.transforms as transforms |
|
from utils_core import flow_viz |
|
import cv2 |
|
|
|
class DDDataset(data.Dataset): |
|
def __init__(self): |
|
super(DDDataset, self).__init__() |
|
def initialize(self, opt): |
|
self.opt = opt |
|
self.dir_txt = opt.datapath |
|
self.paths = [] |
|
in_file = open(self.dir_txt, "r") |
|
k = 0 |
|
list_paths = in_file.readlines() |
|
for line in list_paths: |
|
|
|
flag = False |
|
line = line.strip() |
|
line = line.split() |
|
|
|
|
|
if (not os.path.exists(line[0])): |
|
print(line[0]+" not exists") |
|
continue |
|
if (not os.path.exists(line[1])): |
|
print(line[1]+" not exists") |
|
continue |
|
if (not os.path.exists(line[2])): |
|
print(line[2]+" not exists") |
|
continue |
|
if (not os.path.exists(line[3])): |
|
print(line[3]+" not exists") |
|
continue |
|
|
|
|
|
|
|
|
|
|
|
path_list = [line[0], line[1], line[2], line[3]] |
|
self.paths.append(path_list) |
|
k += 1 |
|
in_file.close() |
|
self.data_size = len(self.paths) |
|
print("num data: ", len(self.paths)) |
|
|
|
def process_data(self, color, mask): |
|
non_zero = mask.nonzero() |
|
bound = 10 |
|
min_x = max(0, non_zero[1].min()-bound) |
|
max_x = min(self.opt.width-1, non_zero[1].max()+bound) |
|
min_y = max(0, non_zero[0].min()-bound) |
|
max_y = min(self.opt.height-1, non_zero[0].max()+bound) |
|
color = color * (mask!=0).astype(float)[:, :, None] |
|
crop_color = color[min_y:max_y, min_x:max_x, :] |
|
crop_color = cv2.resize(np.ascontiguousarray(crop_color), (self.opt.crop_width, self.opt.crop_height), interpolation=cv2.INTER_LINEAR) |
|
crop_params = [[min_x], [max_x], [min_y], [max_y]] |
|
|
|
return crop_color, crop_params |
|
|
|
def __getitem__(self, index): |
|
paths = self.paths[index % self.data_size] |
|
src_color = np.array(Image.open(paths[0])) |
|
src_color = src_color.astype(np.uint8) |
|
raw_src_color = src_color.copy() |
|
src_mask = np.array(Image.open(paths[1])) |
|
src_mask_copy = src_mask.copy() |
|
src_crop_color, src_crop_params = self.process_data(src_color, src_mask) |
|
|
|
|
|
raw_src_color = torch.from_numpy(raw_src_color).permute(2, 0, 1).float() / 255.0 |
|
src_crop_color = torch.from_numpy(src_crop_color).permute(2, 0, 1).float() / 255.0 |
|
|
|
src_mask_copy = (src_mask_copy!=0) |
|
src_mask_copy = torch.tensor(src_mask_copy[np.newaxis, :, :]) |
|
|
|
tar_color = np.array(Image.open(paths[2])) |
|
tar_color = tar_color.astype(np.uint8) |
|
raw_tar_color = tar_color.copy() |
|
tar_mask = np.array(Image.open(paths[3])) |
|
tar_mask_copy = tar_mask.copy() |
|
tar_crop_color, tar_crop_params = self.process_data(tar_color, tar_mask) |
|
|
|
raw_tar_color = torch.from_numpy(raw_tar_color).permute(2, 0, 1).float() / 255.0 |
|
tar_crop_color = torch.from_numpy(tar_crop_color).permute(2, 0, 1).float() / 255.0 |
|
|
|
tar_mask_copy = (tar_mask_copy!=0) |
|
tar_mask_copy = torch.tensor(tar_mask_copy[np.newaxis, :, :]) |
|
|
|
Crop_param = torch.tensor(src_crop_params+tar_crop_params) |
|
|
|
split_ = paths[0].split("/") |
|
path1 = split_[-1][:-4] + "_" + paths[2].split("/")[-1][:-4] +".oflow" |
|
|
|
return {"path_flow":path1, "src_crop_color":src_crop_color, "tar_crop_color":tar_crop_color, "src_color":raw_src_color, "tar_color":raw_tar_color, "src_mask":src_mask_copy, "tar_mask":tar_mask_copy, "Crop_param":Crop_param} |
|
|
|
def __len__(self): |
|
return self.data_size |
|
|
|
def name(self): |
|
return 'DDDataset' |
