def save_images(self,
ids,
kpts,
kpts_gt,
noisy_kpts,
split,
step: int = 0):
"""
Creates a grid of images with gt joints and a grid with predicted joints.
This is a basic function for debugging purposes only.
The grid will be written in the SummaryWriter with name "{prefix}_images" and
"{prefix}_predictions".
Args:
ids (np.array): a list of images ids to be loaded.
kpts (torch.Tensor): a tensor of predicted joints with shape (batch x joints x 2).
kpts_gt (torch.Tensor): a tensor of gt joints with shape (batch x joints x 2).
split (str): split type, train, val or test
step (int): summary_writer step.
Default: 0
Returns:
A pair of images which are built from torchvision.utils.make_grid
"""
# We only need to print first 4 images
if ids.shape[0] > 4:
ids = ids[:4]
kpts = kpts[:4]
kpts_gt = kpts_gt[:4]
# Retrieve Image from dataset
if self.dataset == "itop":
if split == "train":
name = [self.train_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_train
imgs_data = h5py.File(
f"/nas/DriverMonitoring/Datasets/ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['data']
elif split == "val":
name = [self.val_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_train
imgs_data = h5py.File(
f"/nas/DriverMonitoring/Datasets/ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['data']
elif split == "test":
name = [self.test_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_test
imgs_data = h5py.File(
f"/nas/DriverMonitoring/Datasets/ITOP/ITOP_{self.side}_test_depth_map.h5",
'r')['data']
else:
raise ValueError("Split: {} not recognized".format(split))
# Convert 3D annotations to 2D
if self.configer.get("metrics", "kpts_type").lower() == "3d":
kpts = world_to_depth(kpts)
kpts_gt = world_to_depth(kpts_gt)
imgs = list()
for i, name in enumerate(name):
index = int(np.where(np.array(ids_label) == name)[0])
imgs.append(imgs_data[index])
else:
raise NotImplementedError("Dataset: {} not implemented".format(
self.dataset))
imgs_detection = np.array(
[point_on_image(k, el) for k, el in zip(kpts, imgs)])
imgs_gt = np.array(
[point_on_image(k, el) for k, el in zip(kpts_gt, imgs)])
grid_pred = torchvision.utils.make_grid(
torch.from_numpy(imgs_detection).permute(0, 3, 1, 2).float() / 255,
nrow=int(imgs_detection.shape[0]**0.5),
padding=2,
normalize=False)
grid_gt = torchvision.utils.make_grid(
torch.from_numpy(imgs_gt).permute(0, 3, 1, 2).float() / 255,
nrow=int(imgs_gt.shape[0]**0.5),
padding=2,
normalize=False)
split = "validation" if split == "val" else split
self.loss_summary.add_image(split + '_prediction',
grid_pred,
global_step=step)
self.loss_summary.add_image(split + '_gt', grid_gt, global_step=step)
if noisy_kpts is not None:
if self.dataset == "itop":
if noisy_kpts.shape[-1] == 3:
noisy_kpts = world_to_depth(noisy_kpts / 1000)
imgs_noise = np.array(
[point_on_image(k, el) for k, el in zip(noisy_kpts, imgs)])
grid_noise = torchvision.utils.make_grid(
torch.from_numpy(imgs_noise).permute(0, 3, 1, 2).float() / 255,
nrow=int(imgs_detection.shape[0]**0.5),
padding=2,
normalize=False)
self.loss_summary.add_image(split + '_input',
grid_noise,
global_step=step)
return grid_gt, grid_pred
def update_metrics(self,
kpts,
kpts_gt=None,
visible=None,
split: str = "",
ids=None,
oks: OKS = None,
accuracy: Metric_ITOP = None,
additional=None):
"""Funtion to update metrics
Args:
kpts (np.ndarray): Current data to update
kpts_gt (np.ndarray): Ground truth relative to current update
visible (np.ndarray, optional): Visible binary mask
split (np.ndarray, optional): Split type, train test or val
ids (list of str, optional): ids of the current kpts to update
oks (OKS, optional): oks metric variable
accuracy (Metric_ITOP, optional): accuracy metric variable
"""
# Update metrics with new inference
if self.configer.get("metrics", "kpts_type").lower() == "3d":
if self.dataset == "itop":
kpts_2d = world_to_depth(kpts)
kpts_gt_2d = world_to_depth(kpts_gt)
kpts_3d = kpts
kpts_gt_3d = kpts_gt
elif self.configer.get("metrics", "kpts_type").lower() == "2d":
kpts_2d = kpts
kpts_gt_2d = kpts_gt
if self.dataset == "itop":
kpts_3d = np.zeros((kpts.shape[0], 15, 3))
kpts_gt_3d = np.zeros((kpts_gt.shape[0], 15, 3))
if split == "train":
name = [self.train_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_train
imgs = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['data']
elif split == "val":
name = [self.val_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_train
imgs = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['data']
elif split == "test":
name = [self.test_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_test
imgs = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_depth_map.h5",
'r')['data']
else:
raise ValueError("Split: {} not recognized".format(split))
for i, name in enumerate(name):
index = int(np.where(np.array(ids_label) == name)[0])
depth = imgs[index] * 1000
kpts_3d[i] = depth_to_world(kpts[i], depth)
kpts_gt_3d[i] = depth_to_world(kpts_gt[i], depth)
else:
raise NotImplementedError(
"Got 2D kpts metrics not in itop, need to be implemented!")
elif self.configer.get("metrics", "kpts_type").lower() == "rwc":
kpts_2d = kpts[..., :2]
kpts_gt_2d = kpts_gt[..., :2]
kpts_3d = np.zeros((kpts.shape[0], 15, 3))
kpts_gt_3d = np.zeros((kpts_gt.shape[0], 15, 3))
kpts_3d[..., 2] = kpts[..., 2]
kpts_gt_3d[..., 2] = kpts_gt[..., 2]
kpts_3d[..., 0] = (kpts[..., 0] - 160) * 0.0035 * kpts[..., 2]
kpts_3d[..., 1] = -(kpts[..., 1] - 120) * 0.0035 * kpts[..., 2]
kpts_gt_3d[...,
0] = (kpts_gt[..., 0] - 160) * 0.0035 * kpts_gt[..., 2]
kpts_gt_3d[...,
1] = -(kpts_gt[..., 1] - 120) * 0.0035 * kpts_gt[..., 2]
else:
raise NotImplementedError("Not implemented metric type: {}".format(
self.configer.get("metrics", "kpts_type")))
oks.eval(kpts_2d, kpts_gt_2d)
accuracy.eval(kpts_3d, kpts_gt_3d, visible)
if additional is not None:
additional.eval(kpts_3d, kpts_gt_3d, visible)
def __init__(self,
configer,
base_transform=None,
input_transform=None,
label_transform=None,
split="train"):
"""Constructor method for ITOP Dataset class
Args:
configer (Configer): Configer object for current procedure phase (train, test, val)
base_transform (Object, optional): Data augmentation transformation for every data
input_transform (Object, optional): Data augmentation transformation only for input
label_transform (Object, optional): Data augmentation transformation only for labels
split (str, optional): Current procedure phase (train, test, val)
"""
print("Loader started.")
# Setting up useful public variables
self.configer = configer
self.base_transform = base_transform #: Basic data augmentation transform for every data.
self.input_transform = input_transform #: Data augmentation transform for input data.
self.label_transform = label_transform #: Data augmentation transformation only for labels
self.split = split #: Split indicate Test/Train/Val procedure
self.side = self.configer["side"] #: Side for top or side view in ITOP
self.data_path = self.configer["train_dir"] #: Path to data directory
# Setting random state seed for validation and test only
self.mu = self.configer.get("data_aug", "mu")
self.sigma = self.configer.get("data_aug", "sigma")
if self.split == "test":
self.rand_generator = np.random.RandomState(
seed=37) #: Random Generator for adding Noise during test
if self.split == "val":
self.rand_generator = np.random.RandomState(
seed=1295185) #: Random Generator for adding Noise during val
self.type = self.configer.get(
"data", "type"
) #: str: Data description, can be base, depth, voxel or pcloud.
# Loading input pcloud/depth ids
if self.split == 'train' or self.split == 'val':
with open(self.configer.get("data", "kpts_path"), 'rb') as infile:
# ITOP data are saved as Meter, we use everything in mm
data = pickle.load(infile)
if self.type.lower() == 'voxel' or self.type.lower() == 'pcloud':
ids = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_point_cloud.h5",
'r')['id']
elif self.type.lower() == 'depth' or self.type.lower() == 'base':
ids = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['id']
else:
raise NotImplementedError('Data type not supported: {}'.format(
self.type))
if self.split == "train":
self.ids = [
str(el) for el in ids
if not (str(el).startswith("b\'04_")
or str(el).startswith("b\'05_"))
]
else:
self.ids = [
str(el) for el in ids if (str(el).startswith("b\'04_")
or str(el).startswith("b\'05_"))
]
labels = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_labels.h5", 'r')
elif self.split == 'test':
with open(self.configer.get("data", "kpts_path_test"),
'rb') as infile:
data = pickle.load(infile)
if self.type.lower() == 'voxel' or self.type.lower() == 'pcloud':
ids = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_point_cloud.h5",
'r')['id']
elif self.type.lower() == 'depth' or self.type.lower() == 'base':
ids = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_depth_map.h5",
'r')['id']
else:
raise NotImplementedError('Data type not supported: {}'.format(
self.type))
self.ids = [str(el) for el in ids]
labels = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_labels.h5", 'r')
else:
raise NotImplementedError('Split error: {}'.format(self.split))
# Loading full annotations for voxel/pcloud
self.kpts = dict()
self.gt = dict()
self.visible = dict()
tmp = self.ids.copy()
self.ids_str = list(map(
str,
labels['id'])) #: list of str: Ids list for ITOP on test split
if self.type == "base" or self.type == "pcloud":
if self.split in ("train", "val"):
images = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_train_depth_map.h5",
'r')['data']
else:
images = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_depth_map.h5",
'r')['data']
for name in tmp:
index = self.ids_str.index(name)
# If is_valid flag is set to 0,
# you should not use any of the provided human joint locations for the particular frame.
if labels['is_valid'][index] == 0:
# Removing the ids from the correct variable
self.ids.remove(name)
continue
else:
if self.type.lower() in ('voxel', 'pcloud'):
# We process data to work with mm value, normally they are expressed in meters
self.gt[name] = np.array(
labels['real_world_coordinates'][index]) * 1000
elif self.type.lower() == 'depth' or self.type.lower(
) == 'base':
if self.configer.get('metrics',
'kpts_type').lower() == '2d':
self.gt[name] = np.array(
labels['image_coordinates'][index])
elif self.configer.get('metrics',
'kpts_type').lower() == 'rwc':
gt = np.array(labels['real_world_coordinates'][index])
gt[:, :2] = world_to_depth(gt)
gt[:, 2] = gt[:, 2] * 1000
self.gt[name] = gt
else:
# We process data to work with mm value, normally they are expressed in meters
self.gt[name] = np.array(
labels['real_world_coordinates'][index]) * 1000
else:
raise NotImplementedError('Type error: {}'.format(
self.type))
if self.configer.get(
"data",
"from_gt") and self.split.lower() in ("val", "test"):
# retrieve kpts in mm, input unprocessed is expressed in m
if self.type == "base":
in_kpt = np.array(labels['image_coordinates'][index])
self.kpts[name] = depth_to_world(
in_kpt, images[index]) * 1000
# self.kpts[name] = np.array(self.gt[name].copy())
else:
self.kpts[name] = self.gt[name].copy()
if self.configer.get('metrics',
'kpts_type').lower() in ('2d', 'rwc'):
noise = self.rand_generator.normal(
self.mu, self.sigma, (self.kpts[name].shape[0], 2))
self.kpts[name][:, :2] = self.kpts[name][:, :2] + noise
else:
noise = self.rand_generator.normal(
self.mu, self.sigma, self.kpts[name].shape)
if self.type == "base":
if self.sigma < 10:
noise = noise[:, :-1]
tmp = world_to_depth(self.kpts[name])
tmp = tmp + noise
tmp = depth_to_world(tmp, images[index]) * 1000
self.kpts[name] = tmp
else:
self.kpts[name][self.kpts[name][:, 2] != 0] = self.kpts[name][self.kpts[name][:, 2] != 0] \
+ noise[self.kpts[name][:, 2] != 0]
else:
self.kpts[name] = self.kpts[name] + noise
elif self.configer.get("data", "from_gt") and self.split.lower(
) == "train" and self.type == "base":
# Adding here gaussian noise on 2D then calculating RWC on the noise input
in_kpt = np.array(labels['image_coordinates'][index])
noise = np.random.normal(self.mu, self.sigma,
(in_kpt.shape[0], 2))
in_kpt = in_kpt + noise
self.kpts[name] = depth_to_world(in_kpt,
images[index]) * 1000
elif self.configer.get("data", "from_gt") and self.split.lower(
) == "train" and self.type == "pcloud":
self.kpts[name] = self.gt[name]
else:
# retrieve kpts in mm
if self.configer["data", "kpts_type"].lower() == "3d" and data[index].shape[-1] != 3 \
and self.configer.get("data", "from_gt") is False:
self.kpts[name] = depth_to_world(
data[index], images[index] * 1000)
else:
self.kpts[name] = data[index]
self.visible[name] = labels['visible_joints'][index]
if self.split in ("test", "val"):
self.visible[name][self.visible[name] == 0] = 1
if self.type == 'base':
for i, el in enumerate(self.visible[name]):
if el == 0:
self.kpts[name][i] = np.zeros(
(self.kpts[name][i].shape))
self.gt[name][i] = np.zeros((self.gt[name][i].shape))
# Setting dataset size.
self.size = len(self.ids)
print(f"Loaded {self.size} values. Done.")
def save_images(self, ids, kpts):
"""Save image to the directory specified in the configer
Args:
kpts (np.ndarray): Current data to update
ids (list of str, optional): ids of the current kpts to update
"""
if self.dataset == "itop":
name = [self.data_loader.dataset.ids[el] for el in ids]
ids_label = self.ids_test
imgs_data = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_test_depth_map.h5",
'r')['data']
imgs = list()
for i, n in enumerate(name):
index = int(np.where(np.array(ids_label) == n)[0])
imgs.append(imgs_data[index])
# Convert 3D annotations to 2D
if self.configer.get("metrics", "kpts_type").lower() == "3d":
kpts_3d = kpts.copy()
kpts = world_to_depth(kpts)
zaxis = np.zeros((kpts[0].shape[0], 3))
zaxis[:, :-1] = kpts[0]
zaxis[:, -1] = kpts_3d[0][:, -1] / 1000.0
if not os.path.exists(
f"{self.configer.get('data', 'result_dir')}/images/vitruvian"
):
os.makedirs(
f"{self.configer.get('data', 'result_dir')}/images/vitruvian"
)
plot_2D_3D(
zaxis,
imgs[0],
f"{self.configer.get('data', 'result_dir')}/images/vitruvian/{ids[0]}.eps",
stride=1)
elif self.configer.get("metrics", "kpts_type").lower() == "rwc":
kpts = kpts[..., :2]
elif self.configer.get("metrics", "kpts_type").lower() == "2d":
kpts_3d = np.zeros((kpts.shape[0], 15, 3))
if self.dataset == "itop":
for i, n in enumerate(name):
index = int(np.where(np.array(ids_label) == n)[0])
depth = imgs_data[index] * 1000
kpts_3d[i] = depth_to_world(kpts[i], depth)
if not os.path.exists(
f"{self.configer.get('data', 'result_dir')}/images/op_2d_3d"
):
os.makedirs(
f"{self.configer.get('data', 'result_dir')}/images/op_2d_3d"
)
plot_2D_3D(
kpts[0],
imgs[0],
f"{self.configer.get('data', 'result_dir')}/images/op_2d_3d/{ids[0]}.eps",
stride=1)
else:
raise NotImplementedError("Dataset: {} not implemented".format(
self.dataset))
return
imgs_detection = np.array([
point_on_image(k, el, v) for k, el, v in zip(kpts, imgs, visible)
])
if not os.path.exists(
f"{self.configer.get('data', 'result_dir')}/images/patch2d"):
os.makedirs(
f"{self.configer.get('data', 'result_dir')}/images/patch2d")
PATH = f"{self.configer.get('data', 'result_dir')}/images/patch2d"
if not os.path.exists(PATH):
os.makedirs(PATH)
for name, el in zip(ids, imgs_detection):
cv2.imwrite(f"{PATH}/{str(name)}.png", el)
def __getitem__(self, idx):
"""getitem method to iterate over ITOP dataset
Args:
idx (int): Iteration index
Returns:
patches (np.ndarray): Processed patches for patches2D, vitruvian and pcloud of the idx selected value
gt (np.ndarray): Ground truth value for the idx selected value
visible (np.ndarray): Visible mask for the idx selected value
kpts (np.ndarray): Input kpts before extracting the patch for idx selected value
idx (int): Iteration index for current value
"""
name = self.ids[idx]
if self.configer.get("data", "from_gt") and self.split == 'train':
kpts = self.kpts[name].copy()
if self.configer.get('metrics',
'kpts_type').lower() in ('2d', 'rwc'):
noise = np.random.normal(self.mu, self.sigma,
(kpts.shape[0], 2))
kpts[:, :2] = kpts[:, :2] + noise
else:
# If base for vitruvian, noise has already been added at 2D level over kpts public variable
if self.type == 'base':
if self.configer["visible_include"] is None:
for i, el in enumerate(self.visible[name]):
if el == 0:
kpts[i] = np.zeros(self.kpts[name][i].shape)
else:
noise = np.random.normal(self.mu, self.sigma, kpts.shape)
kpts = kpts + noise
else:
kpts = np.array(self.kpts[name])
if 'zaxis' in self.configer['data', 'kpts_path_test']:
kpts = zaxis_to_world_np(kpts)
if self.configer['data', 'kpts_type'] == "2D" and \
('zaxis' in self.configer['data','kpts_path_test'] or '3d' in self.configer['data','kpts_path_test'].lower()):
kpts = world_to_depth(kpts)
gt = self.gt[name].copy()
split = self.split if self.split != "val" else "train"
if self.type.lower() == 'voxel' or self.type.lower() == 'pcloud':
f = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_" + split +
"_point_cloud.h5", 'r')
elif self.type.lower() == 'depth' or self.type.lower() == 'base':
f = h5py.File(
f"{self.data_path}ITOP/ITOP_{self.side}_" + split +
"_depth_map.h5", 'r')
else:
raise NotImplementedError('Data type not supported: {}'.format(
self.type))
index = self.ids_str.index(name)
data = f['data'][index].copy()
# Voxel type
if self.type.lower() == 'voxel':
"""
Todo:
Implement Voxel Extracting procedure
"""
# Pcloud type
elif self.type.lower() == 'pcloud':
data = data * 1000
# Retrieving pcloud dimension
pcloud_dim = self.configer.get("data", "pcloud_dim")
if self.configer["side"] == "top":
N = 5000
else:
N = 2000
patches = np.zeros((15, N + 1, 3))
for i, el in enumerate(kpts):
if self.visible[name][i] == 0 or el[2] == 0:
continue
tmp2 = np.zeros((16, 3))
pcloud_dim_tmp = pcloud_dim
while tmp2.shape[0] <= 128 and pcloud_dim_tmp <= 450:
tmp2 = data[(data[:, 0] > el[0] - pcloud_dim_tmp / 2)
& (data[:, 0] < el[0] + pcloud_dim_tmp / 2) &
(data[:, 1] > el[1] - pcloud_dim_tmp / 2) &
(data[:, 1] < el[1] + pcloud_dim_tmp / 2) &
(data[:, 2] > el[2] - pcloud_dim_tmp / 2) &
(data[:, 2] <
el[2] + pcloud_dim_tmp / 2)].copy()
pcloud_dim_tmp += 50
tmp2 = tmp2[tmp2[:, 2] != 0]
tmp = tmp2.copy()[:N]
# max = 4500 points for every body parts, calculated over train + test + val split
# Need to create cube with 4500 point, correct in the middle, 0 everywhere else
tmp = pointcloud_normalization(tmp)
patches[i][:tmp.shape[0], :] = tmp
patches[i][-1] = len(tmp)
gt = kpts - gt
# 2D depth maps type
elif self.type.lower() == 'depth':
kpts = kpts[:, :2]
mean = mean_depth_patch
std = std_depth_patch
# Retrieving depth patch dimesnion
patch_dim = self.configer.get("data", "patch_dim") // 2
patches = np.zeros((15, patch_dim * 2, patch_dim * 2))
for i, el in enumerate(kpts):
# Slicing for retrieving correct patch
patch = data[int(el[1]) - patch_dim:int(el[1]) + patch_dim,
int(el[0]) - patch_dim:int(el[0]) + patch_dim]
patches[i][:patch.shape[0], :patch.shape[1]] = patch
patches[i][patches[i] == 0] = np.mean(data)
if self.configer.get("data", "from_gt") and self.configer.get(
'metrics', 'kpts_type').lower() == 'rwc':
kpts[i,
2] = data[int(el[1]) if el[1] < 240 else 239,
int(el[0]) if el[0] < 320 else 319] * 1000
patches -= mean
patches /= std
gt = kpts - gt
# Option without real patches, only kpts
# Patches parameters is still used for simplicity
elif self.type.lower() == 'base':
patches = kpts.copy()
if self.configer.get("offset") is True:
gt = kpts - gt
patches[patches[:, 2] != 0] -= mean_base
patches[patches[:, 2] != 0] /= std_base
else:
raise NotImplementedError('Data type not supported: {}'.format(
self.type))
# Data augmentation
if self.base_transform is not None:
patches = self.base_transform(patches)
if self.input_transform is not None:
patches = self.input_transform(patches)
if self.label_transform is not None:
gt = self.label_transform(gt)
return patches.astype(np.float32), gt.astype(
np.float32), self.visible[name].astype(np.float32), kpts.astype(
np.float32), idx