182 lines
6.4 KiB
Python
182 lines
6.4 KiB
Python
import cv2
|
|
import random
|
|
import numpy as np
|
|
|
|
from ..loaders import dl_config as dl_cfg
|
|
from ..loaders import dataloader as dl
|
|
from ..visualize import plotting as plot
|
|
|
|
|
|
def inspect_parser():
|
|
import argparse
|
|
pars = argparse.ArgumentParser(description='Data augmentation and dataset visualization. '
|
|
'Press Q to quit,'
|
|
'N to visualize the next image'
|
|
' and any other key to visualize the next default data.')
|
|
pars.add_argument('database', type=str,
|
|
choices=['wflw', '300wpublic', '300wprivate', 'cofw68', 'merlrav'], help='Database name')
|
|
pars.add_argument('-a', '--anns', type=str, default='train', help='Annotation type: test, train or valid')
|
|
pars.add_argument('-np', '--nopose', action='store_false', default=True, help='Avoid pose generation')
|
|
pars.add_argument('-c', '--clean', action='store_true', help='Process without data augmentation for train')
|
|
pars.add_argument('--shape', nargs='+', type=int, default=[256, 256], help='Image cropped shape (W,H)')
|
|
pars.add_argument('--img', nargs='+', type=int, default=None, help='Select specific image ids')
|
|
return pars.parse_args()
|
|
|
|
|
|
class DatasetInspector:
|
|
|
|
def __init__(self, database, anns_type, data_aug=True, pose=True, image_shape=(256,256)):
|
|
|
|
data_config = dl_cfg.AlignConfig(database, anns_type)
|
|
data_config.image_size = image_shape
|
|
data_config.ftmap_size = image_shape
|
|
data_config.generate_pose = pose
|
|
|
|
if not data_aug:
|
|
data_config.aug_names = []
|
|
|
|
self.data_config = data_config
|
|
dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
|
|
self.dataset = dataset
|
|
self.dataloader = dataloader
|
|
self.colors_dft = {'lnd': (plot.GREEN, plot.RED), 'pose': (plot.BLUE, plot.GREEN, plot.RED)}
|
|
|
|
def show_dataset(self, ids_list=None):
|
|
|
|
if ids_list is None:
|
|
ids = self.get_idx(shuffle=self.data_config.shuffle)
|
|
else:
|
|
ids = ids_list
|
|
|
|
for img_id in ids:
|
|
data_dict = self.dataset[img_id]
|
|
crop_imgs, full_img = self.plot_features(data_dict)
|
|
|
|
# Plot crop
|
|
if 'merge' in crop_imgs.keys():
|
|
crop = crop_imgs['merge']
|
|
else:
|
|
crop = crop_imgs['lnd']
|
|
cv2.imshow('crop', crop)
|
|
|
|
# Plot full
|
|
cv2.imshow('image', full_img['lnd'])
|
|
|
|
key = cv2.waitKey()
|
|
if key == ord('q'):
|
|
break
|
|
|
|
def plot_features(self, data_dict, colors=None):
|
|
|
|
# Init variables
|
|
crop_imgs = {}
|
|
full_imgs = {}
|
|
if colors is None:
|
|
colors = self.colors_dft
|
|
|
|
# Cropped image
|
|
image = data_dict['image']
|
|
landmarks = data_dict['landmarks']
|
|
visible = data_dict['visible']
|
|
if np.any(np.isnan(visible)):
|
|
visible = None
|
|
mask = data_dict['mask_ldm']
|
|
|
|
# Full image
|
|
if 'image_ori' in data_dict.keys():
|
|
image_ori = data_dict['image_ori']
|
|
else:
|
|
image_ori = cv2.imread(data_dict['imgpath'])
|
|
landmarks_ori = data_dict['landmarks_ori']
|
|
visible_ori = data_dict['visible_ori']
|
|
if np.any(np.isnan(visible_ori)):
|
|
visible_ori = None
|
|
mask_ori = data_dict['mask_ldm_ori']
|
|
|
|
# Plot landmarks
|
|
crop_imgs['lnd'] = self._plot_lnd(image, landmarks, visible, mask, colors=colors['lnd'])
|
|
full_imgs['lnd'] = self._plot_lnd(image_ori, landmarks_ori, visible_ori, mask_ori, colors=colors['lnd'])
|
|
|
|
if self.data_config.generate_pose:
|
|
rot, trl, cam_matrix = self._extract_pose(data_dict)
|
|
|
|
# Plot pose
|
|
crop_imgs['pose'] = plot.draw_pose(image, rot, trl, cam_matrix, euler=True, colors=colors['pose'])
|
|
|
|
# Plot merge features
|
|
crop_imgs['merge'] = plot.draw_pose(crop_imgs['lnd'], rot, trl, cam_matrix, euler=True, colors=colors['pose'])
|
|
|
|
return crop_imgs, full_imgs
|
|
|
|
def get_idx(self, shuffle=False):
|
|
ids = list(range(len(self.dataset)))
|
|
if shuffle:
|
|
random.shuffle(ids)
|
|
return ids
|
|
|
|
def reload_dataset(self, data_config=None):
|
|
if data_config is None:
|
|
data_config = self.data_config
|
|
dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
|
|
self.dataset = dataset
|
|
self.dataloader = dataloader
|
|
|
|
def _extract_pose(self, data_dict):
|
|
# Rotation and translation matrix
|
|
pose = data_dict['pose']
|
|
rot = pose[:3]
|
|
trl = pose[3:]
|
|
|
|
# Camera matrix
|
|
cam_matrix = data_dict['cam_matrix']
|
|
|
|
# Check for ground truth anns
|
|
if 'headpose_ori' in data_dict.keys():
|
|
if len(self.data_config.aug_names) == 0:
|
|
print('Image headpose generated by ground truth data')
|
|
pose_ori = data_dict['headpose_ori']
|
|
rot = pose_ori
|
|
|
|
return rot, trl, cam_matrix
|
|
|
|
def _plot_lnd(self, image, landmarks, visible, mask, max_shape_thr=720, colors=None):
|
|
|
|
if colors is None:
|
|
colors = self.colors_dft['lnd']
|
|
|
|
# Full image plots
|
|
W, H, C = image.shape
|
|
|
|
# Original image resize if need it
|
|
if W > max_shape_thr or H > max_shape_thr:
|
|
max_shape = max(W, H)
|
|
scale_factor = max_shape_thr / max_shape
|
|
resize_shape = (int(H * scale_factor), int(W * scale_factor))
|
|
image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask,
|
|
thick_scale=1 / scale_factor, colors=colors)
|
|
image_out = cv2.resize(image_out, resize_shape)
|
|
else:
|
|
image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask, colors=colors)
|
|
|
|
return image_out
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = inspect_parser()
|
|
data_aug = True
|
|
database = args.database
|
|
anns_type = args.anns
|
|
pose = args.nopose
|
|
select_img = args.img
|
|
if args.clean:
|
|
data_aug = False
|
|
|
|
if len(args.shape) != 2:
|
|
raise ValueError('--shape requires two values: width and height. Ej: --shape 256 256')
|
|
else:
|
|
img_shape = tuple(args.shape)
|
|
|
|
visualizer = DatasetInspector(database, anns_type, data_aug=data_aug, pose=pose, image_shape=img_shape)
|
|
visualizer.show_dataset(ids_list=select_img)
|
|
|