def _augment(img, depth, label):
st = lambda x: iaa.Sometimes(0.5, x) # NOQA
augmentations = [
st(iaa.WithChannels([0, 1], iaa.Multiply([1, 1.5]))),
st(
iaa.InColorspace(
'HSV',
children=iaa.WithChannels([1, 2], iaa.Multiply([0.5, 2])),
)),
(iaa.GaussianBlur(sigma=[0, 1])),
iaa.Sometimes(0.9, iaa.Dropout(p=(0, 0.4), name='dropout')),
# iaa.CoarseDropout(p=(0, 0.1), size_percent=0.5, name='dropout'),
iaa.Sometimes(
0.9,
iaa.Affine(
order=1,
cval=0,
scale=1,
translate_px=(-96, 96),
rotate=(-180, 180),
mode='constant',
)),
]
aug = iaa.Sequential(augmentations, random_order=True)
def activator_imgs(images, augmenter, parents, default):
if isinstance(augmenter, iaa.Affine):
augmenter.order = Deterministic(1)
augmenter.cval = Deterministic(0)
return True
def activator_depths(images, augmenter, parents, default):
white_lists = (iaa.Affine, iaa.Sequential, iaa.Sometimes)
if not isinstance(augmenter, white_lists):
return False
if isinstance(augmenter, iaa.Affine):
augmenter.order = Deterministic(1)
augmenter.cval = Deterministic(0)
return True
def activator_lbls(images, augmenter, parents, default):
white_lists = (iaa.Affine, iaa.Sequential, iaa.Sometimes)
if not isinstance(augmenter, white_lists):
return False
if isinstance(augmenter, iaa.Affine):
augmenter.order = Deterministic(0)
augmenter.cval = Deterministic(-1)
return True
aug = aug.to_deterministic()
img = aug.augment_image(img,
hooks=ia.HooksImages(activator=activator_imgs))
depth = aug.augment_image(
depth, hooks=ia.HooksImages(activator=activator_depths))
label = aug.augment_image(
label, hooks=ia.HooksImages(activator=activator_lbls))
return img, depth, label
def get_example(self, i):
files_dir = self._files_dirs[i]
image_file = osp.join(files_dir, 'image.png')
image = skimage.io.imread(image_file)
assert image.dtype == np.uint8
assert image.ndim == 3
label_file = osp.join(files_dir, 'label.png')
with open(label_file, 'r') as f:
label = np.asarray(PIL.Image.open(f)).astype(np.int32)
assert label.dtype == np.int32
assert label.ndim == 2
# Data augmentation
if self.aug:
# 1. Color augmentation
obj_datum = dict(img=image)
random_state = np.random.RandomState()
def st(x):
return iaa.Sometimes(0.3, x)
augs = [
st(iaa.Add([-50, 50], per_channel=True)),
st(
iaa.InColorspace('HSV',
children=iaa.WithChannels(
[1, 2], iaa.Multiply([0.5, 2])))),
st(iaa.GaussianBlur(sigma=[0.0, 1.0])),
st(
iaa.AdditiveGaussianNoise(scale=(0.0, 0.1 * 255),
per_channel=True)),
]
obj_datum = next(
mvtk.aug.augment_object_data([obj_datum],
random_state=random_state,
augmentations=augs))
image = obj_datum['img']
# 2. Geometric augmentation
np.random.seed()
if np.random.uniform() < 0.5:
image = np.fliplr(image)
label = np.fliplr(label)
if np.random.uniform() < 0.5:
image = np.flipud(image)
label = np.flipud(label)
if np.random.uniform() < 0.5:
angle = (np.random.uniform() * 180) - 90
image = self.rotate_image(image, angle, cv2.INTER_LINEAR)
label = self.rotate_image(label, angle, cv2.INTER_NEAREST)
return image, label
def __init__(self,
data_type='all',
random_state=1234,
resize_rate=0.5,
test_size=0.1,
cross_validation=False,
loop=None,
img_aug=False,
with_damage=True,
classification=False):
assert data_type in ('train', 'val', 'all')
# scrape dataset
data_ids = self._get_data_ids()
# split train/val data
if data_type != 'all':
if cross_validation:
random.seed(random_state)
random.shuffle(data_ids)
test_num = int(len(data_ids) * test_size)
start = loop * test_num
end = (loop + 1) * test_num
if len(data_ids) - end > test_num:
ids_train = data_ids[:start] + data_ids[end:]
ids_val = data_ids[start:end]
else:
ids_train = data_ids[:start]
ids_val = data_ids[start:]
else:
ids_train, ids_val = train_test_split(
data_ids, test_size=test_size, random_state=random_state)
if data_type == 'train':
self.data_ids = ids_train
elif data_type == 'val':
self.data_ids = ids_val
else:
self.data_ids = data_ids
self.resize_rate = resize_rate
if img_aug:
self.aug = iaa.Sequential([
iaa.Sometimes(
0.3,
iaa.InColorspace('HSV',
children=iaa.WithChannels(
[1, 2], iaa.Multiply([0.5, 2])))),
iaa.Fliplr(0.5)
])
else:
self.aug = False
if with_damage:
self.failure_label = np.array([
'singlearm_drop',
'singlearm_protrude',
'singlearm_damage',
'dualarm_drop',
'dualarm_protrude',
'dualarm_damage',
])
else:
self.failure_label = np.array([
'singlearm_drop',
'singlearm_protrude',
'dualarm_drop',
'dualarm_protrude',
])
if classification:
with open(self.class_yamlpath, 'r') as f:
self.class_label = yaml.load(f)
else:
self.class_label = None
def _get_example(self, files_dir, idx):
image_file = osp.join(files_dir, 'image.png')
image = skimage.io.imread(image_file)
assert image.dtype == np.uint8
assert image.ndim == 3
depth_file = osp.join(files_dir, 'depth.npz')
depth = np.load(depth_file)['arr_0']
depth[depth == 0] = np.nan
if depth.dtype == np.uint16:
depth = depth.astype(np.float32)
depth /= 1000
depth_keep = ~np.isnan(depth)
depth[depth_keep] = np.maximum(depth[depth_keep], self.min_value)
depth[depth_keep] = np.minimum(depth[depth_keep], self.max_value)
assert depth.dtype == np.float32
assert depth.ndim == 2
label_file = osp.join(files_dir, 'label.png')
with open(label_file, 'r') as f:
label = np.asarray(PIL.Image.open(f)).astype(np.int32)
assert label.dtype == np.int32
assert label.ndim == 2
depth_gt_file = osp.join(files_dir, 'depth_gt.npz')
depth_gt = np.load(depth_gt_file)['arr_0']
depth_gt[depth_gt == 0] = np.nan
if depth_gt.dtype == np.uint16:
depth_gt = depth_gt.astype(np.float32)
depth /= 1000
depth_gt_keep = ~np.isnan(depth_gt)
depth_gt[depth_gt_keep] = np.maximum(
depth_gt[depth_gt_keep], self.min_value)
depth_gt[depth_gt_keep] = np.minimum(
depth_gt[depth_gt_keep], self.max_value)
assert depth_gt.dtype == np.float32
assert depth_gt.ndim == 2
# Data augmentation
if self.aug:
# 1. Color augmentation
obj_datum = dict(img=image)
random_state = np.random.RandomState()
def st(x):
return iaa.Sometimes(0.3, x)
augs = [
st(iaa.Add([-50, 50], per_channel=True)),
st(iaa.InColorspace(
'HSV', children=iaa.WithChannels(
[1, 2], iaa.Multiply([0.5, 2])))),
st(iaa.GaussianBlur(sigma=[0.0, 1.0])),
st(iaa.AdditiveGaussianNoise(
scale=(0.0, 0.1 * 255), per_channel=True)),
]
obj_datum = next(mvtk.aug.augment_object_data(
[obj_datum], random_state=random_state, augmentations=augs))
image = obj_datum['img']
# 2. Depth noise
np.random.seed()
if np.random.uniform() < 0.3:
noise_rate = np.random.uniform() * 0.25 + 0.05
depth[
np.random.rand(depth.shape[0], depth.shape[1]) < noise_rate
] = np.nan
# 3. Geometric augmentation
if np.random.uniform() < 0.5:
image = np.fliplr(image)
depth = np.fliplr(depth)
label = np.fliplr(label)
depth_gt = np.fliplr(depth_gt)
if np.random.uniform() < 0.5:
image = np.flipud(image)
depth = np.flipud(depth)
label = np.flipud(label)
depth_gt = np.flipud(depth_gt)
if np.random.uniform() < 0.5:
angle = (np.random.uniform() * 180) - 90
image = self.rotate_image(image, angle, cv2.INTER_LINEAR)
depth = self.rotate_depth_image(depth, angle, cv2.INTER_LINEAR)
label = self.rotate_image(label, angle, cv2.INTER_NEAREST)
depth_gt = self.rotate_depth_image(
depth_gt, angle, cv2.INTER_LINEAR)
return image, depth, label, depth_gt, idx