def __init__(
self,
images_dir: str,
mats_dir: str,
landmarks_dir: str = None,
transform: bool = False,
):
self.images, self.volumes, self.landmarks = [], [], []
self.transform = transform
if transform:
self.tf_flip = data_transform.Flip()
self.tf_rotate = data_transform.Rotation()
self.tf_translate = data_transform.Translation()
self.tf_scale = data_transform.Scale()
for i in os.listdir(images_dir):
name = i.split(".")[0]
self.images += [os.path.join(images_dir, name + ".jpg")]
self.volumes += [os.path.join(mats_dir, name + ".mat")]
if landmarks_dir:
self.landmarks += [os.path.join(landmarks_dir, name + ".mat")]
assert len(self.images) == len(self.volumes)
def __init__(self,
images_dir: str,
mats_dir: str,
landmarks_dir: str = None,
transform: bool = False):
self.images, self.volumes, self.landmarks = [], [], []
self.transform = transform
if transform:
self.tf_flip = data_transform.Flip()
self.tf_rotate = data_transform.Rotation()
self.tf_translate = data_transform.Translation()
self.tf_scale = data_transform.Scale()
for i in os.listdir(images_dir):
ext = os.path.splitext(i)[1]
if ext == '.jpg' or ext == '.jpeg':
self.images += [os.path.join(images_dir, i)]
for j in os.listdir(mats_dir):
ext = os.path.splitext(j)[1]
if ext == '.mat':
self.volumes += [os.path.join(mats_dir, j)]
if landmarks_dir:
for j in os.listdir(landmarks_dir):
ext = os.path.splitext(j)[1]
if ext == '.mat':
self.landmarks += [os.path.join(landmarks_dir, i)]
assert len(self.images) == len(self.volumes)
def __getitem__(self, idx):
# read input image
if self.input_format == 'img':
rgb_name = os.path.join(self.root_dir,
self.rgbd_frame.iloc[idx, 0])
with open(rgb_name, 'rb') as fRgb:
rgb_image = Image.open(rgb_name).convert('RGB')
depth_name = os.path.join(self.root_dir,
self.rgbd_frame.iloc[idx, 1])
with open(depth_name, 'rb') as fDepth:
depth_image = Image.open(depth_name)
# read input hdf5
elif self.input_format == 'hdf5':
file_name = os.path.join(self.root_dir,
self.rgbd_frame.iloc[idx, 0])
rgb_h5, depth_h5 = self.load_h5(file_name)
rgb_image = Image.fromarray(rgb_h5, mode='RGB')
depth_image = Image.fromarray(depth_h5.astype('float32'), mode='F')
else:
print('error: the input format is not supported now!')
return None
_s = np.random.uniform(1.0, 1.5)
s = np.int(240*_s)
degree = np.random.uniform(-5.0, 5.0)
if self.split == 'train':
tRgb = data_transform.Compose([transforms.Resize(s),
data_transform.Rotation(degree),
transforms.ColorJitter(brightness = 0.4, contrast = 0.4, saturation = 0.4),
# data_transform.Lighting(0.1, imagenet_eigval, imagenet_eigvec)])
transforms.CenterCrop((228, 304)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
transforms.ToPILImage()])
tDepth = data_transform.Compose([transforms.Resize(s),
data_transform.Rotation(degree),
transforms.CenterCrop((228, 304))])
rgb_image = tRgb(rgb_image)
depth_image = tDepth(depth_image)
if np.random.uniform()<0.5:
rgb_image = rgb_image.transpose(Image.FLIP_LEFT_RIGHT)
depth_image = depth_image.transpose(Image.FLIP_LEFT_RIGHT)
rgb_image = transforms.ToTensor()(rgb_image)
if self.input_format == 'img':
depth_image = transforms.ToTensor()(depth_image)
else:
depth_image = data_transform.ToTensor()(depth_image)
depth_image = depth_image.div(_s)
sparse_image = self.createSparseDepthImage(depth_image, self.n_sample)
rgbd_image = torch.cat((rgb_image, sparse_image), 0)
sample = {'rgbd': rgbd_image, 'depth': depth_image}
elif self.split == 'val':
tRgb = data_transform.Compose([transforms.Resize(240),
transforms.CenterCrop((228, 304)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
transforms.ToPILImage()])
tDepth = data_transform.Compose([transforms.Resize(240),
transforms.CenterCrop((228, 304))])
rgb_raw = tDepth(rgb_image)
rgb_image = tRgb(rgb_image)
depth_image = tDepth(depth_image)
rgb_image = transforms.ToTensor()(rgb_image)
rgb_raw = transforms.ToTensor()(rgb_raw)
if self.input_format == 'img':
depth_image = transforms.ToTensor()(depth_image)
else:
depth_image = data_transform.ToTensor()(depth_image)
sparse_image = self.createSparseDepthImage(depth_image, self.n_sample)
rgbd_image = torch.cat((rgb_image, sparse_image), 0)
sample = {'rgbd': rgbd_image, 'depth': depth_image, 'raw_rgb': rgb_raw }
return sample