def __call__(self, input_dict : {str : Union[Image.Image, np.ndarray]}) -> dict:
image, label = input_dict[tag_image], input_dict[tag_label]
image = np.array(image)
label = np.array(label)
# size measure
y_max = image.shape[0]
x_max = image.shape[1]
# np.ndarray -> imgaug.augmentables.segmaps.SegmentationMapsOnImage
label = SegmentationMapsOnImage(label, shape=image.shape)
# augmentation
zoomset = iaa.OneOf([
iaa.Identity(), # do nothing
iaa.Affine(scale=self.outscale), # zoom out
RandomCrop(y_max, x_max).cut() # zoom in
])
image, label = zoomset(image=image, segmentation_maps=label)
image, label = self.transformSet(image=image, segmentation_maps=label)
# imgaug.augmentables.segmaps.SegmentationMapsOnImage -> np.ndarray
label = label.get_arr()
return {tag_image : image,
tag_label : label}
def apply_imgaug_to_imgs(self, rgb, depth, obj_mask, aff_mask,
obj_part_mask):
rgb, depth = np.array(rgb), np.array(depth)
aff_mask, obj_mask, obj_part_mask = np.array(aff_mask), np.array(
obj_mask), np.array(obj_part_mask)
H, W, C = rgb.shape[0], rgb.shape[1], rgb.shape[2]
concat_img = np.zeros(shape=(H, W, C + 1))
concat_img[:, :, :C] = rgb
concat_img[:, :, -1] = depth
concat_img = np.array(concat_img, dtype=np.uint8)
concat_mask = np.zeros(shape=(H, W, 3))
concat_mask[:, :, 0] = aff_mask
concat_mask[:, :, 1] = obj_mask
concat_mask[:, :, 2] = obj_part_mask
concat_mask = np.array(concat_mask, dtype=np.uint8)
segmap = SegmentationMapsOnImage(concat_mask,
shape=np.array(rgb).shape)
aug_concat_img, segmap = self.affine(image=concat_img,
segmentation_maps=segmap)
aug_concat_mask = segmap.get_arr()
rgb = aug_concat_img[:, :, :C]
depth = aug_concat_img[:, :, -1]
aff_mask = aug_concat_mask[:, :, 0]
obj_mask = aug_concat_mask[:, :, 1]
obj_part_mask = aug_concat_mask[:, :, 2]
rgb = self.colour_aug(image=rgb)
depth = self.depth_aug(image=depth)
rgb = np.array(rgb, dtype=np.uint8)
aff_mask = np.array(aff_mask, dtype=np.uint8)
obj_mask = np.array(obj_mask, dtype=np.uint8)
obj_part_mask = np.array(obj_part_mask, dtype=np.uint8)
depth = np.array(depth, dtype=np.uint8)
# Uncomment to check resulting images with augmentation.
# helper_utils.print_class_labels(aff_mask)
# helper_utils.print_class_labels(obj_mask)
# helper_utils.print_class_labels(obj_part_mask)
# cv2.imshow('rgb', cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB))
# color_aff_mask = ycb_video_dataset_utils.colorize_aff_mask(aff_mask)
# cv2.imshow('color_aff_mask', cv2.cvtColor(color_aff_mask, cv2.COLOR_BGR2RGB))
# color_obj_mask = ycb_video_dataset_utils.colorize_obj_mask(obj_mask)
# cv2.imshow('color_obj_mask', cv2.cvtColor(color_obj_mask, cv2.COLOR_BGR2RGB))
# color_obj_part_mask = ycb_video_dataset_utils.convert_obj_part_mask_to_obj_mask(obj_part_mask)
# color_obj_part_mask = ycb_video_dataset_utils.colorize_obj_mask(color_obj_part_mask)
# cv2.imshow('color_obj_part_mask', cv2.cvtColor(color_obj_part_mask, cv2.COLOR_BGR2RGB))
# cv2.waitKey(0)
return rgb, depth, obj_mask, aff_mask, obj_part_mask
def __call__(self, img, segmap, softmax):
segmap = SegmentationMapsOnImage(segmap, shape=segmap.shape)
softmax = HeatmapsOnImage(softmax, shape=softmax.shape)
assert img.shape[:2] == segmap.shape[:2] == softmax.shape[:2]
for a in self.augmentations:
img, segmap, softmax = a(img, segmap, softmax)
return img, segmap.get_arr(), softmax.get_arr()
def __getitem__(self, i):
idx = self.ids[i]
mask_file = glob(self.masks_dir + idx + self.mask_suffix + '.*')
img_file = glob(self.imgs_dir + idx + '.*')
depth_file = glob(self.depth_dir + idx + self.depth_suffix + '.*')
assert len(mask_file) == 1, \
f'Either no mask or multiple masks found for the ID {idx}: {mask_file}'
assert len(img_file) == 1, \
f'Either no image or multiple images found for the ID {idx}: {img_file}'
assert len(depth_file) == 1, \
f'Either no image or multiple images found for the ID {idx}: {depth_file}'
mask = Image.open(mask_file[0])
img = Image.open(img_file[0])
depth = Image.open(depth_file[0])
### 3-channel depth images
depth = Image.fromarray(
skimage.color.gray2rgb(np.array(depth, dtype=np.uint8)))
assert img.size == mask.size and depth.size == mask.size, \
f'Image and mask {idx} should be the same size, but are {img.size}, {depth.size}' \
f' and {mask.size}'
################################
### TODO: IMGAUG
################################
if self.take_center_crop:
img = self.crop(img, self.crop_w, self.crop_w, is_img=True)
depth = self.crop(depth, self.crop_w, self.crop_w, is_img=True)
mask = self.crop(mask, self.crop_w, self.crop_w, is_img=False)
if self.apply_imgaug:
img, depth, mask = np.array(img), np.array(depth), np.array(mask)
segmap = SegmentationMapsOnImage(mask, shape=np.array(img).shape)
image, segmap = self.affine(image=img, segmentation_maps=segmap)
depth = self.affine(image=depth)
mask = segmap.get_arr()
img = self.colour_aug(image=image)
depth = self.depth_aug(image=depth)
img, depth, mask = Image.fromarray(img), Image.fromarray(
depth), Image.fromarray(mask)
img = self.preprocess(img, self.scale, is_img=True)
depth = self.preprocess(depth, self.scale, is_img=True)
mask = self.preprocess(mask, self.scale, is_img=False)
return {
'image': torch.from_numpy(img).type(torch.FloatTensor),
'depth': torch.from_numpy(depth).type(torch.FloatTensor),
'mask': torch.from_numpy(mask).type(torch.FloatTensor)
}
def forward(self, data):
aug1 = self.aug_non_spatial.to_deterministic()
aug2 = self.aug_spatial.to_deterministic()
data['image'] = aug1(image=data['image'])
if 'label' in data.keys() and data['label'] is not None:
segmap = SegmentationMapsOnImage(data['lable'],
shape=data['image'].shape)
data['image'], segmap = aug2(image=data['image'],
segmentation_maps=segmap)
data['label'] = segmap.get_arr()
else:
data['image'] = aug2(image=data['image'])
return data
def __getitem__(self, idx):
ia.seed(idx + 1)
pt_idx, env_idx = self.idx_list[idx]
if self.args.data_mode == '2D':
probe_img = self.env_dict[env_idx]['img'][pt_idx].astype(np.float)
probe_mask = self.env_dict[env_idx]['mask'][pt_idx].astype(
np.float)
probe_img = np.expand_dims(probe_img, axis=-1)
probe_mask = np.expand_dims(probe_mask, axis=-1)
elif self.args.data_mode == '2.5D':
img_stack_list = []
img_stack_list.append(self.env_dict[env_idx]['img'][pt_idx].astype(
np.float))
step = int((self.args.slices - 1) / 2)
for i in range(step):
s_idx = max(pt_idx - sum([i for i in range(i + 1)]), 0)
e_idx = min(pt_idx + sum([i for i in range(i + 1)]),
len(self.env_dict[env_idx]['img']) - 1)
img_stack_list.insert(
0, self.env_dict[env_idx]['img'][s_idx].astype(np.float))
img_stack_list.insert(
-1, self.env_dict[env_idx]['img'][e_idx].astype(np.float))
probe_img = np.stack(img_stack_list, axis=-1)
probe_mask = self.env_dict[env_idx]['mask'][pt_idx].astype(
np.float)
probe_mask = np.expand_dims(probe_mask, axis=-1)
else:
print(self.args.data_mode + " is not implemented.")
raise NotImplementedError
probe_img, probe_mask = center_crop(probe_img, probe_mask,
self.args.crop_size)
probe_mask = probe_mask.astype(np.int32)
# augmentation
if self.augmentation:
seg_map = SegmentationMapsOnImage(probe_mask,
shape=probe_img.shape)
aug_affine = iaa.Affine(scale=(0.9, 1.1),
translate_percent=(-0.05, 0.05),
rotate=(-360, 360),
shear=(-20, 20),
mode='edge')
probe_img, seg_map = aug_affine(image=probe_img,
segmentation_maps=seg_map)
probe_mask = seg_map.get_arr()
sample = {'img': probe_img, 'mask': probe_mask}
return sample
def generate_train(self, batch_size, augmentation=True):
assert self.get_nb_train() >= (
self.timestep * batch_size
), "Pas assez d'images de train pour ce batch_size et ce timestep, il en faut au moins {} et il y en a {}".format(
self.timestep * batch_size, self.get_nb_train())
i = 0
nb = self.get_nb_train()
while i <= (nb // batch_size):
batch_input = np.zeros(
(batch_size, self.timestep, self.input_shape[0],
self.input_shape[1], 3),
dtype=np.float32)
batch_target = np.zeros(
(batch_size, self.input_shape[0], self.input_shape[1], 1),
dtype=np.uint8)
for j in range(batch_size):
images_to_use = self.im_train[i * batch_size + j]
if augmentation:
aug_det = self.augmentation.to_deterministic()
for index, im in enumerate(images_to_use):
img = Image.open(im)
img = img.resize(
(self.input_shape[1], self.input_shape[0]))
img_array = np.array(img)
if augmentation:
img_array = aug_det.augment_image(img_array)
batch_input[j, index] = ((img_array - 127.5) /
127.5).astype(np.float32)
gt = Image.open(self.gt_train[i * batch_size + j])
gt = gt.resize((self.input_shape[1], self.input_shape[0]))
gt = np.array(gt, dtype=np.uint8)
gt = np.expand_dims(gt, axis=-1)
if gt.dtype != np.bool:
gt[(gt < 255) & (gt > 1)] = 0
gt[gt > 0] = 1
gt = gt.astype(np.bool)
if augmentation:
gt = SegmentationMapsOnImage(gt,
shape=(img_array.shape[0],
img_array.shape[1]))
gt = aug_det.augment_segmentation_maps(gt)
gt = gt.get_arr()
batch_target[j] = gt
i += 1
if i >= (nb // batch_size):
i = 0
yield batch_input, batch_target
def __getitem__(self, item):
image = Image.open(self.image_list[item]).convert("RGB")
image_ori = Image.open(self.seg_list[item]).convert("RGB")
if self.augment:
image = self.augment(image)
image_ori = np.array(image_ori)
image = np.array(image)
image_ori = SegmentationMapsOnImage(image_ori, shape=image_ori.shape)
if self.transform:
image, image_ori = self.transform(image=image, segmentation_maps=image_ori)
image_ori = image_ori.get_arr()
image = (np.array(image, dtype=np.float32) / 255.0).transpose((2, 0, 1))
image_ori = (np.array(image_ori, dtype=np.float32) / 255.0).transpose((2, 0, 1))
image = (image - 0.5) * 2
image_ori = (image_ori - 0.5) * 2
return image, image_ori
def augstore(self, src:dict, dst_base:str,
dataname_extension='.tiff', labelname_extension='.tif',
identifier=None):
os.makedirs(dst_base, exist_ok=True)
os.makedirs(os.path.join(dst_base, label_folder_name), exist_ok=True)
# get image
image = src[tag_image] # PIL.Image.Image
label = src[tag_label] # PIL.Image.Image
name = src[tag_name] # str
# PIL -> numpy
image = np.array(image)
label = np.array(label)
# size measure
y_max = image.shape[0]
x_max = image.shape[1]
# np.ndarray -> imgaug.augmentables.segmaps.SegmentationMapsOnImage
label = SegmentationMapsOnImage(label, shape=label.shape)
# augmentation
zoomset = iaa.OneOf([
iaa.Identity(), # do nothing
iaa.Affine(scale=self.outscale), # zoom out
RandomCrop(y_max, x_max).cut() # zoom in
])
image, label = zoomset(image=image, segmentation_maps=label)
image, label = self.transformSet(image=image, segmentation_maps=label)
# imgaug.augmentables.segmaps.SegmentationMapsOnImage -> np.ndarray
label = label.get_arr()
if not identifier == None:
name = name + '_' + str(identifier)
# numpy -> PIL.Image.Image
image = Image.fromarray(image)
label = Image.fromarray(label)
image.save(os.path.join(dst_base, name + dataname_extension))
label.save(os.path.join(dst_base, label_folder_name, name + labelname_extension))
return {tag_image : image,
tag_label : label,
tag_name : name}
def load_mask(self, image_id):
info = self.image_info[image_id]
imagePath = info["path"]
maskPath = info["mask_path"]
augmentation = info["augmentation_params"]
image = skimage.io.imread(imagePath)
image = kutils.RCNNConvertInputImage(image)
mask = skimage.io.imread(maskPath)
if mask.ndim > 2:
if mask.shape[0] < mask.shape[2]:
mask = mask[0, :, :]
else:
mask = mask[:, :, 0]
segmap = SegmentationMapsOnImage(mask, shape=image.shape)
if augmentation is not None:
segmap = augmentation(segmentation_maps=segmap)
mask = segmap.get_arr()
maskIndexMax = numpy.max(mask)
newMaskIndices = numpy.zeros([maskIndexMax], dtype=numpy.uint32)
for y in range(mask.shape[0]):
for x in range(mask.shape[1]):
index = int(mask[y, x]) - 1
if index >= 0:
newMaskIndices[index] = 1
count = 0
for i in range(maskIndexMax):
if newMaskIndices[i] > 0:
newMaskIndices[i] = count
count += 1
masks = numpy.zeros([mask.shape[0], mask.shape[1], count],
dtype=numpy.uint8)
for y in range(mask.shape[0]):
for x in range(mask.shape[1]):
index = int(mask[y, x]) - 1
if index >= 0:
masks[y, x, newMaskIndices[index]] = 1
#assign class id 1 to all masks
class_ids = numpy.array([1 for _ in range(count)])
return masks, class_ids.astype(numpy.int32)
def process_image_segmentation(image, segmap, input_w, input_h, output_w,
output_h, augment):
# resize the image to standard size
if (input_w and input_h) or augment:
segmap = SegmentationMapsOnImage(segmap, shape=image.shape)
if (input_w and input_h):
# Rescale image and segmaps
image = ia.imresize_single_image(image, (input_w, input_h))
segmap = segmap.resize((output_w, output_h),
interpolation="nearest")
if augment:
aug_pipe = _create_augment_pipeline()
image, segmap = aug_pipe(image=image, segmentation_maps=segmap)
return image, segmap.get_arr()
def demo(self, img_dir, out_dir):
os.makedirs(out_dir, exist_ok=True)
img_paths = sorted(glob(osp.join(img_dir, '*.png')), key=osp.getmtime)
to_tensor = transforms.ToTensor()
resize = iaa.Resize({
"height": self.opt.dataset.rgb_res[0],
"width": self.opt.dataset.rgb_res[1]
})
with torch.no_grad():
self.model.eval()
for idx, img_path in enumerate(img_paths):
print(img_path)
# prepare input
rgb_img = cv2.imread(img_path)
rgb_img = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2RGB)
resize_to_orig = iaa.Resize({
"height": rgb_img.shape[0],
"width": rgb_img.shape[1]
})
rgb_img = resize(image=rgb_img)
rgb_img_ts = to_tensor(rgb_img)
rgb_img_ts = rgb_img_ts.unsqueeze(0).to(self.device)
# forward
pred_mask = self.model(rgb_img_ts)
# transform ouput
pred_mask = pred_mask[0].cpu() # (2,256,256), float
pred_mask = torch.argmax(pred_mask, 0).numpy().astype('uint8')
pred_segmap = SegmentationMapsOnImage(pred_mask,
shape=rgb_img.shape)
rgb_img, pred_segmap = resize_to_orig(
image=rgb_img, segmentation_maps=pred_segmap)
# write results
pred_mask = pred_segmap.get_arr().astype('uint8')
pred_mask = (pred_mask * 255).astype('uint8')
cv2.imwrite(osp.join(out_dir,
img_path.split('/')[-1]), pred_mask)
cells = []
cells.append(rgb_img)
cells.append(pred_segmap.draw_on_image(rgb_img)[0])
cells.append(pred_segmap.draw(size=rgb_img.shape[:2])[0])
grid_image = ia.draw_grid(cells, rows=1, cols=3)
grid_image_path = osp.join(
out_dir,
img_path.split('/')[-1].replace('.png', '_grid.png'))
cv2.imwrite(grid_image_path, grid_image[:, :, ::-1])
def generator(index):
path = paths[index]
# open images
img = imread(path, resize=resize)
# open all masks based on the image path
mask = np.zeros((*img.shape[:2], self.num_classes), dtype='uint8')
path_mask = path.replace('image.jpg', 'mask')
for pm in glob(path_mask + '*.jpg'):
# get the class for the mask
clas = int(pm.replace('.jpg', '').split('_')[-1])
msk = imread(pm, resize=resize)[..., 0]
# adjust mask
msk = cv2.threshold(msk, 0, 1,
cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
mask[..., clas] = msk
mask = SegmentationMapsOnImage(mask, shape=img.shape)
# augmentation
if augment:
img, mask = self.augment(img, mask)
mask = mask.get_arr()
img = self.norm(img)
# draw contours
contours = np.zeros_like(mask)
for i in range(mask.shape[-1]):
base = np.zeros_like(img)
cnts = cv2.findContours(mask[..., i], cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)[0]
contours[..., i] = cv2.drawContours(base, cnts, -1, (1, 1, 1),
2)[..., 0]
yield np.float32(img), {
'mask_output': np.float32(mask),
'contour_output': np.float32(contours)
}
def chapter_examples_segmentation_maps_array():
import imageio
import numpy as np
import imgaug as ia
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
# Load an example image (uint8, 128x128x3).
image = ia.quokka(size=(128, 128), extract="square")
# Create an example segmentation map (int32, 128x128).
# Here, we arbitrarily place some squares on the image.
# Class 0 is the background class.
segmap = np.zeros((128, 128, 1), dtype=np.int32)
segmap[28:71, 35:85, 0] = 1
segmap[10:25, 30:45, 0] = 2
segmap[10:25, 70:85, 0] = 3
segmap[10:110, 5:10, 0] = 4
segmap[118:123, 10:110, 0] = 5
segmap1 = SegmentationMapsOnImage(segmap, shape=image.shape)
# Read out the segmentation map's array, change it and create a new
# segmentation map
arr = segmap1.get_arr()
arr[10:110, 5:10, 0] = 5
segmap2 = ia.SegmentationMapsOnImage(arr, shape=image.shape)
# Draw three columns: (1) original image, (2) original image with
# unaltered segmentation map on top, (3) original image with altered
# segmentation map on top
cells = [
image,
segmap1.draw_on_image(image)[0],
segmap2.draw_on_image(image)[0]
]
# Convert cells to grid image and save.
grid_image = ia.draw_grid(cells, cols=3)
# imageio.imwrite("example_segmaps_array.jpg", grid_image)
save("examples_segmentation_maps", "array.jpg", grid_image, quality=90)
def apply_imgaug_to_imgs(self, rgb, depth, mask):
rgb, depth, mask = np.array(rgb), np.array(depth), np.array(mask)
H, W, C = rgb.shape[0], rgb.shape[1], rgb.shape[2]
concat_img = np.zeros(shape=(H, W, C + 1))
concat_img[:, :, :C] = rgb
concat_img[:, :, -1] = depth
concat_img = np.array(concat_img, dtype=np.uint8)
concat_mask = np.array(mask, dtype=np.uint8)
segmap = SegmentationMapsOnImage(concat_mask,
shape=np.array(rgb).shape)
aug_concat_img, segmap = self.affine(image=concat_img,
segmentation_maps=segmap)
aug_concat_mask = segmap.get_arr()
rgb = aug_concat_img[:, :, :C]
depth = aug_concat_img[:, :, -1]
mask = aug_concat_mask
rgb = self.colour_aug(image=rgb)
depth = self.depth_aug(image=depth)
rgb = np.array(rgb, dtype=np.uint8)
mask = np.array(mask, dtype=np.uint8)
# Uncomment to check resulting images with augmentation.
# dataset_utils.print_class_labels(mask)
# cv2.imshow('rgb', cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB))
# color_mask = umd_dataset_utils.colorize_aff_mask(mask)
# cv2.imshow('color_mask', cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))
# cv2.waitKey(0)
return rgb, depth, mask
def __getitem__(self, idx):
img_path = os.path.join(self.imgs_root,
self.all_in_one[idx][0]) # 0th one= image
mask_path = os.path.join(self.masks_root,
self.all_in_one[idx][1]) # 1st one = mask
grid_size = self.all_in_one[idx][2] # 2nd one = grid size
#print("img path=", img_path)
#print("mask_pth", mask_path)
#print("grid size=", grid_size)
img = Image.open(img_path)
mask = Image.open(mask_path)
grid_encode = generate_checkerboard(256, 256, grid_size) #[:, :, 0]
#print("finished grid encode")
# print()
# resizing
img = img.resize((256, 256), Image.NEAREST)
mask = mask.resize((256, 256), Image.NEAREST)
# covert to numpy
img = np.array(img)
mask = np.array(mask) #
# clean mask values (this is done to remove small values in mask images)
mask = (
mask > 128
) * 255 # if mask value > 128, set it to 255 (this will remove 254, 253 values and 1,2 etc)
#========================================
# add new augmentations to img and mask
#========================================
# convert array to SegmentationMapsOnImage instance
mask = mask.astype(np.int8)
#print("mask type:", type(mask))
segmap = SegmentationMapsOnImage(mask, shape=img.shape)
# First, augment image.
image_aug = self.aug_img(image=img)
img = image_aug
# Second, augment segmentation map.
# We convert to uint8 as that dtype has usually best support and hence is safest to use.
mask_arr_aug = self.aug_mask(image=segmap.get_arr().astype(np.uint8))
mask_aug = SegmentationMapsOnImage(mask_arr_aug, shape=segmap.shape)
mask = mask_aug.get_arr()
# get tiled ground truth for the augmented mask
mask = get_tiled_ground_truth(mask, grid_size)
#mask = mask[:, :, 0]
# To tensor
# mask = np.array(mask)
#mask = torch.from_numpy(mask)
# TO tensor
# img = np.asarray(img)
# img = torch.from_numpy(img)
# mask to BB
#bb_img, boxes = mask_to_bb(mask)
# convert everything into a torch.Tensor
#boxes = torch.as_tensor(boxes, dtype=torch.float32)
#image_id = torch.tensor([idx])
#area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
# suppose all instances are not crowd
#iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)
#labels = torch.ones((len(boxes),), dtype=torch.int64)
# area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
#target = {}
# target["bb_img"] = bb_img
#target["boxes"] = boxes
#target["labels"] = labels
# target["masks"] = mask
#target["image_id"] = image_id
#target["area"] = area
#target["iscrowd"] = iscrowd
if self.transforms is not None:
# img, target= self.transforms(img, target)
img = self.transforms(img)
mask = self.transforms(mask)
grid_encode = self.transforms(grid_encode)
#t = torch.Tensor([0.5]) # threshold
#mask = (mask > t).float() * 1
# print("transform applited..!")
# print("Image:", img.type)
#print("mask:", target["masks"].type)
# print("boxes;", target["boxes"].type)
# print("image id:", target["image_id"].type)
#print("labels;", target["labels"].type)
# print("iscrowd:",target["iscrowd"].type)
# print(img.type)
return {"img": img, "grid_encode": grid_encode, "mask": mask}
def test_Flipud():
reseed()
base_img = np.array([[0, 0, 1], [0, 0, 1], [0, 1, 1]], dtype=np.uint8)
base_img = base_img[:, :, np.newaxis]
base_img_flipped = np.array([[0, 1, 1], [0, 0, 1], [0, 0, 1]],
dtype=np.uint8)
base_img_flipped = base_img_flipped[:, :, np.newaxis]
images = np.array([base_img])
images_flipped = np.array([base_img_flipped])
keypoints = [
ia.KeypointsOnImage([
ia.Keypoint(x=0, y=0),
ia.Keypoint(x=1, y=1),
ia.Keypoint(x=2, y=2)
],
shape=base_img.shape)
]
keypoints_flipped = [
ia.KeypointsOnImage([
ia.Keypoint(x=0, y=3 - 0),
ia.Keypoint(x=1, y=3 - 1),
ia.Keypoint(x=2, y=3 - 2)
],
shape=base_img.shape)
]
polygons = [
ia.PolygonsOnImage([ia.Polygon([(0, 0), (2, 0), (2, 2)])],
shape=base_img.shape)
]
polygons_flipped = [
ia.PolygonsOnImage([ia.Polygon([(0, 3 - 0), (2, 3 - 0), (2, 3 - 2)])],
shape=base_img.shape)
]
# 0% chance of flip
aug = iaa.Flipud(0)
aug_det = aug.to_deterministic()
for _ in sm.xrange(10):
observed = aug.augment_images(images)
expected = images
assert np.array_equal(observed, expected)
observed = aug_det.augment_images(images)
expected = images
assert np.array_equal(observed, expected)
observed = aug.augment_keypoints(keypoints)
expected = keypoints
assert keypoints_equal(observed, expected)
observed = aug_det.augment_keypoints(keypoints)
expected = keypoints
assert keypoints_equal(observed, expected)
for aug_ in [aug, aug_det]:
observed = aug_.augment_polygons(polygons)
assert len(observed) == 1
assert len(observed[0].polygons) == 1
assert observed[0].shape == polygons[0].shape
assert observed[0].polygons[0].exterior_almost_equals(
polygons[0].polygons[0])
assert observed[0].polygons[0].is_valid
# 0% chance of flip, heatmaps
aug = iaa.Flipud(0)
heatmaps = HeatmapsOnImage(np.float32([
[0, 0.5, 0.75],
[0, 0.5, 0.75],
[0.75, 0.75, 0.75],
]),
shape=(3, 3, 3))
observed = aug.augment_heatmaps([heatmaps])[0]
expected = heatmaps.get_arr()
assert observed.shape == heatmaps.shape
assert heatmaps.min_value - 1e-6 < observed.min_value < heatmaps.min_value + 1e-6
assert heatmaps.max_value - 1e-6 < observed.max_value < heatmaps.max_value + 1e-6
assert np.array_equal(observed.get_arr(), expected)
# 0% chance of flip, segmaps
aug = iaa.Flipud(0)
segmaps = SegmentationMapsOnImage(np.int32([
[0, 1, 2],
[0, 1, 2],
[2, 2, 2],
]),
shape=(3, 3, 3))
observed = aug.augment_segmentation_maps([segmaps])[0]
expected = segmaps.get_arr()
assert observed.shape == segmaps.shape
assert np.array_equal(observed.get_arr(), expected)
# 100% chance of flip
aug = iaa.Flipud(1.0)
aug_det = aug.to_deterministic()
for _ in sm.xrange(10):
observed = aug.augment_images(images)
expected = images_flipped
assert np.array_equal(observed, expected)
observed = aug_det.augment_images(images)
expected = images_flipped
assert np.array_equal(observed, expected)
observed = aug.augment_keypoints(keypoints)
expected = keypoints_flipped
assert keypoints_equal(observed, expected)
observed = aug_det.augment_keypoints(keypoints)
expected = keypoints_flipped
assert keypoints_equal(observed, expected)
for aug_ in [aug, aug_det]:
observed = aug_.augment_polygons(polygons)
assert len(observed) == 1
assert len(observed[0].polygons) == 1
assert observed[0].shape == polygons[0].shape
assert observed[0].polygons[0].exterior_almost_equals(
polygons_flipped[0].polygons[0])
assert observed[0].polygons[0].is_valid
# 100% chance of flip, heatmaps
aug = iaa.Flipud(1.0)
heatmaps = ia.HeatmapsOnImage(np.float32([
[0, 0.5, 0.75],
[0, 0.5, 0.75],
[0.75, 0.75, 0.75],
]),
shape=(3, 3, 3))
observed = aug.augment_heatmaps([heatmaps])[0]
expected = np.flipud(heatmaps.get_arr())
assert observed.shape == heatmaps.shape
assert heatmaps.min_value - 1e-6 < observed.min_value < heatmaps.min_value + 1e-6
assert heatmaps.max_value - 1e-6 < observed.max_value < heatmaps.max_value + 1e-6
assert np.array_equal(observed.get_arr(), expected)
# 100% chance of flip, segmaps
aug = iaa.Flipud(1.0)
segmaps = SegmentationMapsOnImage(np.int32([
[0, 1, 2],
[0, 1, 2],
[2, 2, 2],
]),
shape=(3, 3, 3))
observed = aug.augment_segmentation_maps([segmaps])[0]
expected = np.flipud(segmaps.get_arr())
assert observed.shape == segmaps.shape
assert np.array_equal(observed.get_arr(), expected)
# 50% chance of flip
aug = iaa.Flipud(0.5)
aug_det = aug.to_deterministic()
nb_iterations = 1000
nb_images_flipped = 0
nb_images_flipped_det = 0
nb_keypoints_flipped = 0
nb_keypoints_flipped_det = 0
nb_polygons_flipped = 0
nb_polygons_flipped_det = 0
for _ in sm.xrange(nb_iterations):
observed = aug.augment_images(images)
if np.array_equal(observed, images_flipped):
nb_images_flipped += 1
observed = aug_det.augment_images(images)
if np.array_equal(observed, images_flipped):
nb_images_flipped_det += 1
observed = aug.augment_keypoints(keypoints)
if keypoints_equal(observed, keypoints_flipped):
nb_keypoints_flipped += 1
observed = aug_det.augment_keypoints(keypoints)
if keypoints_equal(observed, keypoints_flipped):
nb_keypoints_flipped_det += 1
observed = aug.augment_polygons(polygons)
if observed[0].polygons[0].exterior_almost_equals(
polygons_flipped[0].polygons[0]):
nb_polygons_flipped += 1
observed = aug_det.augment_polygons(polygons)
if observed[0].polygons[0].exterior_almost_equals(
polygons_flipped[0].polygons[0]):
nb_polygons_flipped_det += 1
assert int(nb_iterations * 0.3) <= nb_images_flipped <= int(
nb_iterations * 0.7)
assert int(nb_iterations * 0.3) <= nb_keypoints_flipped <= int(
nb_iterations * 0.7)
assert int(nb_iterations * 0.3) <= nb_polygons_flipped <= int(
nb_iterations * 0.7)
assert nb_images_flipped_det in [0, nb_iterations]
assert nb_keypoints_flipped_det in [0, nb_iterations]
assert nb_polygons_flipped_det in [0, nb_iterations]
# 50% chance of flipped, multiple images, list as input
images_multi = [base_img, base_img]
aug = iaa.Flipud(0.5)
aug_det = aug.to_deterministic()
nb_iterations = 1000
nb_flipped_by_pos = [0] * len(images_multi)
nb_flipped_by_pos_det = [0] * len(images_multi)
for _ in sm.xrange(nb_iterations):
observed = aug.augment_images(images_multi)
for i in sm.xrange(len(images_multi)):
if np.array_equal(observed[i], base_img_flipped):
nb_flipped_by_pos[i] += 1
observed = aug_det.augment_images(images_multi)
for i in sm.xrange(len(images_multi)):
if np.array_equal(observed[i], base_img_flipped):
nb_flipped_by_pos_det[i] += 1
for val in nb_flipped_by_pos:
assert int(nb_iterations * 0.3) <= val <= int(nb_iterations * 0.7)
for val in nb_flipped_by_pos_det:
assert val in [0, nb_iterations]
# test StochasticParameter as p
aug = iaa.Flipud(p=iap.Choice([0, 1], p=[0.7, 0.3]))
seen = [0, 0]
for _ in sm.xrange(1000):
observed = aug.augment_image(base_img)
if np.array_equal(observed, base_img):
seen[0] += 1
elif np.array_equal(observed, base_img_flipped):
seen[1] += 1
else:
assert False
assert 700 - 75 < seen[0] < 700 + 75
assert 300 - 75 < seen[1] < 300 + 75
# test exceptions for wrong parameter types
got_exception = False
try:
_ = iaa.Flipud(p="test")
except Exception:
got_exception = True
assert got_exception
# test get_parameters()
aug = iaa.Flipud(p=0.5)
params = aug.get_parameters()
assert isinstance(params[0], iap.Binomial)
assert isinstance(params[0].p, iap.Deterministic)
assert 0.5 - 1e-4 < params[0].p.value < 0.5 + 1e-4
###################
# test other dtypes
###################
aug = iaa.Flipud(1.0)
image = np.zeros((3, 3), dtype=bool)
image[0, 0] = True
expected = np.zeros((3, 3), dtype=bool)
expected[2, 0] = True
image_aug = aug.augment_image(image)
assert image_aug.dtype.type == image.dtype.type
assert np.all(image_aug == expected)
for dtype in [
np.uint8, np.uint16, np.uint32, np.uint64, np.int8, np.int32,
np.int64
]:
min_value, center_value, max_value = iadt.get_value_range_of_dtype(
dtype)
value = max_value
image = np.zeros((3, 3), dtype=dtype)
image[0, 0] = value
expected = np.zeros((3, 3), dtype=dtype)
expected[2, 0] = value
image_aug = aug.augment_image(image)
assert image_aug.dtype.type == dtype
assert np.array_equal(image_aug, expected)
for dtype, value in zip([np.float16, np.float32, np.float64, np.float128],
[5000, 1000**2, 1000**3, 1000**4]):
atol = 1e-9 * max_value if dtype != np.float16 else 1e-3 * max_value
image = np.zeros((3, 3), dtype=dtype)
image[0, 0] = value
expected = np.zeros((3, 3), dtype=dtype)
expected[2, 0] = value
image_aug = aug.augment_image(image)
assert image_aug.dtype.type == dtype
assert np.allclose(image_aug, expected, atol=atol)
def __call__(self, img, mask):
mask = SegmentationMapsOnImage(mask, shape=img.shape)
img, mask = self.aug(image=img, segmentation_maps=mask)
return img, mask.get_arr()
