def random_morphological_transform(label):
""" Randomly erode/dilate the label
@param label: a [H x W] numpy array of {0, 1}
"""
num_tries = 0
valid_transform = False
while not valid_transform:
if num_tries >= cfg.TRAIN.max_augmentation_tries:
print('Morph: Exhausted number of augmentation tries...')
return label
# Sample whether we do erosion or dilation, and kernel size for that
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(label)
sidelength = np.mean([x_max - x_min, y_max - y_min])
morphology_kernel_size = 0
num_ksize_tries = 0
while morphology_kernel_size == 0:
if num_ksize_tries >= 50: # 50 tries for this
print(
'Morph: Exhausted number of augmentation tries... Sidelength: {sidelength}'
)
return label
dilation_percentage = np.random.beta(
cfg.TRAIN.label_dilation_alpha, cfg.TRAIN.label_dilation_beta)
morphology_kernel_size = int(
round(sidelength * dilation_percentage))
num_ksize_tries += 1
iterations = np.random.randint(1, cfg.TRAIN.morphology_max_iters + 1)
# Erode/dilate the mask
kernel = cv2.getStructuringElement(
cv2.MORPH_ELLIPSE,
(morphology_kernel_size, morphology_kernel_size))
if np.random.rand() < 0.5:
morphed_label = cv2.erode(label, kernel, iterations=iterations)
else:
morphed_label = cv2.dilate(label, kernel, iterations=iterations)
# Make sure there the mass is reasonable
if (np.count_nonzero(morphed_label) / morphed_label.size > 0.001) and \
(np.count_nonzero(morphed_label) / morphed_label.size < 0.98):
valid_transform = True
num_tries += 1
return morphed_label
def crop_rois(rgb, initial_masks, depth):
N, H, W = initial_masks.shape
crop_size = cfg.TRAIN.SYN_CROP_SIZE
padding_percentage = 0.25
mask_ids = torch.unique(initial_masks[0])
if mask_ids[0] == 0:
mask_ids = mask_ids[1:]
num = mask_ids.shape[0]
rgb_crops = torch.zeros((num, 3, crop_size, crop_size), device=cfg.device)
rois = torch.zeros((num, 4), device=cfg.device)
mask_crops = torch.zeros((num, crop_size, crop_size), device=cfg.device)
if depth is not None:
depth_crops = torch.zeros((num, 3, crop_size, crop_size), device=cfg.device)
else:
depth_crops = None
for index, mask_id in enumerate(mask_ids):
mask = (initial_masks[0] == mask_id).float() # Shape: [H x W]
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(mask)
x_padding = int(torch.round((x_max - x_min).float() * padding_percentage).item())
y_padding = int(torch.round((y_max - y_min).float() * padding_percentage).item())
# pad and be careful of boundaries
x_min = max(x_min - x_padding, 0)
x_max = min(x_max + x_padding, W-1)
y_min = max(y_min - y_padding, 0)
y_max = min(y_max + y_padding, H-1)
rois[index, 0] = x_min
rois[index, 1] = y_min
rois[index, 2] = x_max
rois[index, 3] = y_max
# crop
rgb_crop = rgb[0, :, y_min:y_max+1, x_min:x_max+1] # [3 x crop_H x crop_W]
mask_crop = mask[y_min:y_max+1, x_min:x_max+1] # [crop_H x crop_W]
if depth is not None:
depth_crop = depth[0, :, y_min:y_max+1, x_min:x_max+1] # [3 x crop_H x crop_W]
# resize
new_size = (crop_size, crop_size)
rgb_crop = F.upsample_bilinear(rgb_crop.unsqueeze(0), new_size)[0] # Shape: [3 x new_H x new_W]
rgb_crops[index] = rgb_crop
mask_crop = F.upsample_nearest(mask_crop.unsqueeze(0).unsqueeze(0), new_size)[0,0] # Shape: [new_H, new_W]
mask_crops[index] = mask_crop
if depth is not None:
depth_crop = F.upsample_bilinear(depth_crop.unsqueeze(0), new_size)[0] # Shape: [3 x new_H x new_W]
depth_crops[index] = depth_crop
return rgb_crops, mask_crops, rois, depth_crops
def random_cut(label):
""" Randomly cut part of mask
@param label: a [H x W] numpy array of {0, 1}
"""
H, W = label.shape
num_tries = 0
valid_transform = False
while not valid_transform:
if num_tries >= cfg.TRAIN.max_augmentation_tries:
print('Cut: Exhausted number of augmentation tries...')
return label
cut_label = label.copy()
# Sample cut percentage
cut_percentage = np.random.uniform(cfg.TRAIN.cut_percentage_min,
cfg.TRAIN.cut_percentage_max)
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(label)
if np.random.rand() < 0.5: # choose width
sidelength = x_max - x_min
if np.random.rand() < 0.5: # from the left
x = int(round(cut_percentage * sidelength)) + x_min
cut_label[y_min:y_max + 1, x_min:x] = 0
else: # from the right
x = x_max - int(round(cut_percentage * sidelength))
cut_label[y_min:y_max + 1, x:x_max + 1] = 0
else: # choose height
sidelength = y_max - y_min
if np.random.rand() < 0.5: # from the top
y = int(round(cut_percentage * sidelength)) + y_min
cut_label[y_min:y, x_min:x_max + 1] = 0
else: # from the bottom
y = y_max - int(round(cut_percentage * sidelength))
cut_label[y:y_max + 1, x_min:x_max + 1] = 0
# Make sure the mass is reasonable
if (np.count_nonzero(cut_label) / cut_label.size > 0.001) and \
(np.count_nonzero(cut_label) / cut_label.size < 0.98):
valid_transform = True
num_tries += 1
return cut_label
def random_translation(label):
""" Randomly translate mask
@param label: a [H x W] numpy array of {0, 1}
"""
num_tries = 0
valid_transform = False
while not valid_transform:
if num_tries >= cfg.TRAIN.max_augmentation_tries:
print('Translate: Exhausted number of augmentation tries...')
return label
# Get tight bbox of mask
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(label)
sidelength = max(x_max - x_min, y_max - y_min)
# sample translation pixels
translation_percentage = np.random.beta(cfg.TRAIN.translation_alpha,
cfg.TRAIN.translation_beta)
translation_percentage = max(translation_percentage,
cfg.TRAIN.translation_percentage_min)
translation_max = int(round(translation_percentage * sidelength))
translation_max = max(translation_max,
1) # To make sure things don't error out
tx = np.random.randint(-translation_max, translation_max)
ty = np.random.randint(-translation_max, translation_max)
translated_label = translate(label,
tx,
ty,
interpolation=cv2.INTER_NEAREST)
# Make sure the mass is reasonable
if (np.count_nonzero(translated_label) / translated_label.size > 0.001) and \
(np.count_nonzero(translated_label) / translated_label.size < 0.98):
valid_transform = True
num_tries += 1
return translated_label
def pad_crop_resize(self, img, label, depth):
""" Crop the image around the label mask, then resize to 224x224
"""
H, W, _ = img.shape
# sample an object to crop
K = np.max(label)
while True:
if K > 0:
idx = np.random.randint(1, K + 1)
else:
idx = 0
foreground = (label == idx).astype(np.float32)
# get tight box around label/morphed label
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(foreground)
cx = (x_min + x_max) / 2
cy = (y_min + y_max) / 2
# make bbox square
x_delta = x_max - x_min
y_delta = y_max - y_min
if x_delta > y_delta:
y_min = cy - x_delta / 2
y_max = cy + x_delta / 2
else:
x_min = cx - y_delta / 2
x_max = cx + y_delta / 2
sidelength = x_max - x_min
padding_percentage = np.random.uniform(
cfg.TRAIN.min_padding_percentage,
cfg.TRAIN.max_padding_percentage)
padding = int(round(sidelength * padding_percentage))
if padding == 0:
padding = 25
# Pad and be careful of boundaries
x_min = max(int(x_min - padding), 0)
x_max = min(int(x_max + padding), W - 1)
y_min = max(int(y_min - padding), 0)
y_max = min(int(y_max + padding), H - 1)
# crop
if (y_min == y_max) or (x_min == x_max):
continue
img_crop = img[y_min:y_max + 1, x_min:x_max + 1]
label_crop = label[y_min:y_max + 1, x_min:x_max + 1]
roi = [x_min, y_min, x_max, y_max]
if depth is not None:
depth_crop = depth[y_min:y_max + 1, x_min:x_max + 1]
break
# resize
s = cfg.TRAIN.SYN_CROP_SIZE
img_crop = cv2.resize(img_crop, (s, s))
label_crop = cv2.resize(label_crop, (s, s),
interpolation=cv2.INTER_NEAREST)
if depth is not None:
depth_crop = cv2.resize(depth_crop, (s, s),
interpolation=cv2.INTER_NEAREST)
else:
depth_crop = None
return img_crop, label_crop, depth_crop
def random_add(label):
""" Randomly add part of mask
@param label: a [H x W] numpy array of {0, 1}
"""
H, W = label.shape
num_tries = 0
valid_transform = False
while not valid_transform:
if num_tries >= cfg.TRAIN.max_augmentation_tries:
print('Add: Exhausted number of augmentation tries...')
return label
added_label = label.copy()
# Sample add percentage
add_percentage = np.random.uniform(cfg.TRAIN.add_percentage_min,
cfg.TRAIN.add_percentage_max)
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(label)
# Sample translation from center
translation_percentage_x = np.random.uniform(0, 2 * add_percentage)
tx = int(round((x_max - x_min) * translation_percentage_x))
translation_percentage_y = np.random.uniform(0, 2 * add_percentage)
ty = int(round((y_max - y_min) * translation_percentage_y))
if np.random.rand() < 0.5: # choose x direction
sidelength = x_max - x_min
ty = np.random.choice(
[-1, 1]
) * ty # mask will be moved to the left/right. up/down doesn't matter
if np.random.rand() < 0.5: # mask copied from the left.
x = int(round(add_percentage * sidelength)) + x_min
try:
temp = added_label[y_min + ty:y_max + 1 + ty,
x_min - tx:x - tx]
added_label[y_min + ty:y_max + 1 + ty,
x_min - tx:x - tx] = np.logical_or(
temp, added_label[y_min:y_max + 1,
x_min:x])
except ValueError as e: # indices were out of bounds
num_tries += 1
continue
else: # mask copied from the right
x = x_max - int(round(add_percentage * sidelength))
try:
temp = added_label[y_min + ty:y_max + 1 + ty,
x + tx:x_max + 1 + tx]
added_label[y_min + ty:y_max + 1 + ty,
x + tx:x_max + 1 + tx] = np.logical_or(
temp, added_label[y_min:y_max + 1,
x:x_max + 1])
except ValueError as e: # indices were out of bounds
num_tries += 1
continue
else: # choose y direction
sidelength = y_max - y_min
tx = np.random.choice([
-1, 1
]) * tx # mask will be moved up/down. lef/right doesn't matter
if np.random.rand() < 0.5: # from the top
y = int(round(add_percentage * sidelength)) + y_min
try:
temp = added_label[y_min - ty:y - ty,
x_min + tx:x_max + 1 + tx]
added_label[y_min - ty:y - ty,
x_min + tx:x_max + 1 + tx] = np.logical_or(
temp, added_label[y_min:y,
x_min:x_max + 1])
except ValueError as e: # indices were out of bounds
num_tries += 1
continue
else: # from the bottom
y = y_max - int(round(add_percentage * sidelength))
try:
temp = added_label[y + ty:y_max + 1 + ty,
x_min + tx:x_max + 1 + tx]
added_label[y + ty:y_max + 1 + ty,
x_min + tx:x_max + 1 + tx] = np.logical_or(
temp, added_label[y:y_max + 1,
x_min:x_max + 1])
except ValueError as e: # indices were out of bounds
num_tries += 1
continue
# Make sure the mass is reasonable
if (np.count_nonzero(added_label) / added_label.size > 0.001) and \
(np.count_nonzero(added_label) / added_label.size < 0.98):
valid_transform = True
num_tries += 1
return added_label
def random_ellipses(label):
""" Randomly add/drop a few ellipses in the mask
This is adapted from the DexNet 2.0 code.
Their code: https://github.com/BerkeleyAutomation/gqcnn/blob/75040b552f6f7fb264c27d427b404756729b5e88/gqcnn/sgd_optimizer.py
@param label: a [H x W] numpy array of {0, 1}
"""
H, W = label.shape
num_tries = 0
valid_transform = False
while not valid_transform:
if num_tries >= cfg.TRAIN.max_augmentation_tries:
print('Ellipse: Exhausted number of augmentation tries...')
return label
new_label = label.copy()
# Sample number of ellipses to include/dropout
num_ellipses = np.random.poisson(cfg.TRAIN.num_ellipses_mean)
# Sample ellipse centers by sampling from Gaussian at object center
pixel_indices = util_.build_matrix_of_indices(H, W)
h_idx, w_idx = np.where(new_label)
mu = np.mean(pixel_indices[h_idx, w_idx, :],
axis=0) # Shape: [2]. y_center, x_center
sigma = 2 * np.cov(pixel_indices[h_idx, w_idx, :].T) # Shape: [2 x 2]
if np.any(np.isnan(mu)) or np.any(np.isnan(sigma)):
print(mu, sigma, h_idx, w_idx)
ellipse_centers = np.random.multivariate_normal(
mu, sigma, size=num_ellipses) # Shape: [num_ellipses x 2]
ellipse_centers = np.round(ellipse_centers).astype(int)
# Sample ellipse radii and angles
x_min, y_min, x_max, y_max = util_.mask_to_tight_box(new_label)
scale_factor = max(
x_max - x_min, y_max -
y_min) * cfg.TRAIN.ellipse_size_percentage # Mean of gamma r.v.
x_radii = np.random.gamma(cfg.TRAIN.ellipse_gamma_base_shape *
scale_factor,
cfg.TRAIN.ellipse_gamma_base_scale,
size=num_ellipses)
y_radii = np.random.gamma(cfg.TRAIN.ellipse_gamma_base_shape *
scale_factor,
cfg.TRAIN.ellipse_gamma_base_scale,
size=num_ellipses)
angles = np.random.randint(0, 360, size=num_ellipses)
# Dropout ellipses
for i in range(num_ellipses):
center = ellipse_centers[i, :]
x_radius = np.round(x_radii[i]).astype(int)
y_radius = np.round(y_radii[i]).astype(int)
angle = angles[i]
# include or dropout the ellipse
mask = np.zeros_like(new_label)
mask = cv2.ellipse(mask,
tuple(center[::-1]), (x_radius, y_radius),
angle=angle,
startAngle=0,
endAngle=360,
color=1,
thickness=-1)
if np.random.rand() < 0.5:
new_label[mask == 1] = 0 # Drop out ellipse
else:
new_label[mask == 1] = 1 # Add ellipse
# Make sure the mass is reasonable
if (np.count_nonzero(new_label) / new_label.size > 0.001) and \
(np.count_nonzero(new_label) / new_label.size < 0.98):
valid_transform = True
num_tries += 1
return new_label