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 normalize_segmentation_maps(inputs, shapes=None):
# TODO get rid of this deferred import
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
shapes = _preprocess_shapes(shapes)
ntype = estimate_segmaps_norm_type(inputs)
_assert_exactly_n_shapes_partial = functools.partial(
_assert_exactly_n_shapes,
from_ntype=ntype,
to_ntype="List[SegmentationMapsOnImage]",
shapes=shapes)
if ntype == "None":
return None
elif ntype in ["array[int]", "array[uint]", "array[bool]"]:
_assert_single_array_ndim(inputs, 4, "(N,H,W,#SegmapsPerImage)",
"SegmentationMapsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
if ntype == "array[bool]":
return [
SegmentationMapsOnImage(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
return [
SegmentationMapsOnImage(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
elif ntype == "SegmentationMapsOnImage":
return [inputs]
elif ntype == "iterable[empty]":
return None
elif ntype in [
"iterable-array[int]", "iterable-array[uint]",
"iterable-array[bool]"
]:
_assert_many_arrays_ndim(inputs, 3, "(H,W,#SegmapsPerImage)",
"SegmentationMapsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
if ntype == "iterable-array[bool]":
return [
SegmentationMapsOnImage(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
return [
SegmentationMapsOnImage(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
else:
assert ntype == "iterable-SegmentationMapsOnImage", (
"Got unknown normalization type '%s'." % (ntype, ))
return inputs # len allowed to differ from len of images
def __call__(self, image, masks=None):
det_augmenter = self.frame_shift_augmenter.to_deterministic()
if masks:
masks_np, is_binary_mask = [], []
for mask in masks:
if isinstance(mask, BinaryMask):
masks_np.append(mask.tensor().numpy().astype(np.bool))
is_binary_mask.append(True)
elif isinstance(mask, np.ndarray):
masks_np.append(mask.astype(np.bool))
is_binary_mask.append(False)
else:
raise ValueError("Invalid mask type: {}".format(
type(mask)))
num_instances = len(masks_np)
masks_np = SegmentationMapsOnImage(self.condense_masks(masks_np),
shape=image.shape[:2])
# to keep track of which points in the augmented image are padded zeros, we augment an additional all-ones
# array. The problem is that imgaug will apply one augmentation to the image and associated mask, but a
# different augmentation to this array. To prevent this, we manually seed the rng in imgaug before both
# function calls to ensure that the same augmentation is applied to both.
seed = int(datetime.now().strftime('%M%S%f')[-8:])
imgaug.seed(seed)
aug_image, aug_masks = det_augmenter(
image=self.basic_augmenter(image=image),
segmentation_maps=masks_np)
imgaug.seed(seed)
invalid_pts_mask = det_augmenter(
image=np.ones(image.shape[:2] + (1, ), np.uint8)).squeeze(2)
aug_masks = self.expand_masks(aug_masks.get_arr(), num_instances)
aug_masks = [
BinaryMask(mask) if is_bm else mask
for mask, is_bm in zip(aug_masks, is_binary_mask)
]
return aug_image, aug_masks, invalid_pts_mask == 0
else:
masks = [
SegmentationMapsOnImage(np.ones(image.shape[:2], np.bool),
shape=image.shape[:2])
]
aug_image, invalid_pts_mask = det_augmenter(
image=image, segmentation_maps=masks)
return aug_image, invalid_pts_mask.get_arr() == 0
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 test_augment_segmaps__kernel_size_is_two__no_keep_size(self):
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
arr = np.int32([
[0, 1, 2],
[1, 2, 3]
])
segmaps = SegmentationMapsOnImage(arr, shape=(6, 6, 3))
aug = self.augmenter(2, keep_size=False)
segmaps_aug = aug.augment_segmentation_maps(segmaps)
expected = segmaps.resize((1, 2))
assert segmaps_aug.shape == (3, 3, 3)
assert segmaps_aug.arr.shape == (1, 2, 1)
assert np.allclose(segmaps_aug.arr, expected.arr)
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 __call__(self, img_mask):
img, mask = np.asarray(img_mask[0]), np.asarray(img_mask[1])
segmask = SegmentationMapsOnImage(mask, shape=img.shape)
crack_img, crack_mask = self.aug(image=img, segmentation_maps=segmask)
crack_mask = crack_mask.arr[:,:,0] # convert to array
return Image.fromarray(crack_img.astype('uint8')), Image.fromarray(crack_mask.astype('uint8'))
def data_aug(img, gt):
# get segmentation map object and corresponding image
segmap = SegmentationMapsOnImage(gt, shape=gt.shape)
# define transformation rule
seq = iaa.Sequential([
iaa.Fliplr(p=0.5),
iaa.Flipud(p=0.5),
iaa.Multiply((0.8, 1.2), per_channel=0.2),
iaa.ContrastNormalization((0.75, 1.5), per_channel=True),
# iaa.Affine(
# scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
# translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
# rotate=(-25, 25),
# shear=(-8, 8))
# iaa.Sharpen()
], random_order=True)
# augmentation
aug_img, aug_gt = seq(image=img, segmentation_maps=segmap)
aug_gt = 255*aug_gt.get_arr()
# save augmented images
# img_name = "./data/training/augmented_images/images/" + str(i + 1) + ".png"
# gt_name = "./data/training/augmented_images/groundtruth/" + str(i + 1) + ".png"
# cv2.imwrite(img_name, aug_img)
# cv2.imwrite(gt_name, 255*aug_gt.get_arr())
return aug_img, aug_gt
def show_sample_image(self):
idx = np.random.choice(len(self.image_paths), 1)[0]
img_path = self.image_paths[idx]
seg_path = self.label_paths[idx]
img = self.load_image(img_path)
seg = self.load_image(seg_path)
# segmap = SegmentationMapsOnImage(seg[..., 0], shape=img.shape)
segmap = SegmentationMapsOnImage(seg, shape=img.shape)
img_aug, seg_aug = self.augment_func(image=img,
segmentation_maps=segmap)
ia.imshow(np.hstack([img[..., ::-1], seg]))
# seg[..., 0] = seg_aug.get_arr_int()
seg = seg_aug.get_arr_int()
# print(np.nonzero(seg[..., 0]))
# non_index = np.nonzero(seg[..., 0])
# for i in range(len(non_index[0])):
# print(seg[..., 0][non_index[0][i], non_index[1][i]])
ia.imshow(np.hstack([img_aug[..., ::-1], seg]))
seg = seg / 255
seg[seg < 0.5] = 0
seg[seg >= 0.5] = 1
ia.imshow(np.hstack([img_aug[..., ::-1], seg]))
def valid_transform(image, segmentation_maps=None):
if segmentation_maps is None:
return image
segmentation_maps = SegmentationMapsOnImage(segmentation_maps,
shape=image.shape)
return image, segmentation_maps.arr
def _imgaug(self, images, masks):
'''
Perform data augmentation by imgaug library
Args:
images : a batch of images.
masks : a batch of masks.
Returns:
A batch of images and a batch of masks.
'''
images = np.array(images, dtype=np.uint8)
masks = np.array(masks, dtype=np.int32)
seg_maps = []
for i in range(len(masks)):
seg_maps.append(
SegmentationMapsOnImage(masks[i], shape=images[i].shape))
images, seg_maps = self.aug_seq(images=images,
segmentation_maps=seg_maps)
for i in range(len(masks)):
masks[i] = seg_maps[i].get_arr()
masks = np.array(masks, dtype=np.int32)
return images, masks
def main():
'''
Load original images and masks, augment them and save as images.
'''
for dataset in DATASET:
data = np.load(os.path.join('Data', 'Image', 'Labelbox',
dataset.lower() + '_segmented.npz'))
images = data['images']
masks = data['masks'].astype(bool)
image_shape = images.shape[1:3]
images_augmented, masks_augmented = ([], [])
for image, mask in progress(zip(images, masks)):
segmap = SegmentationMapsOnImage(mask, shape=image_shape)
image_aug, segmap_aug = augment_data(SEQ, image, segmap, N_AUGMENT)
for img_aug, sg_aug in zip(image_aug, segmap_aug):
sg_map = sg_aug.draw(size=image_shape)[0]
sg_map = ((sg_map[..., 0] != 0)*255).astype(np.uint8)
images_augmented.append(img_aug)
masks_augmented.append(sg_map)
images_augmented = np.array(images_augmented)
masks_augmented = np.array(masks_augmented)
np.savez_compressed(os.path.join('Data', 'Image', 'Augmented', dataset.lower() +
'_augmented'), images=images_augmented, masks=masks_augmented)
def __call__(self, imgs, annos, use_image):
if use_image:
seq = iaa.Sequential([
iaa.Crop(percent=(0.0, 0.1), keep_size=True),
iaa.Affine(scale=(0.9, 1.1), shear=(-15, 15), rotate=(-25, 25))
])
else:
seq = iaa.Sequential([
iaa.Crop(percent=(0.0, 0.1), keep_size=True),
iaa.Affine(scale=(0.95, 1.05), shear=(-10, 10), rotate=(-15, 15))
])
seq = seq.to_deterministic()
num = len(imgs)
for idx in range(1, num):
img = imgs[idx]
anno = annos[idx]
max_obj = anno.shape[2]-1
anno = convert_one_hot(anno, max_obj)
segmap = SegmentationMapsOnImage(anno, shape=img.shape)
img_aug, segmap_aug = seq(image=img, segmentation_maps=segmap)
imgs[idx] = img_aug
annos[idx] = convert_mask(segmap_aug.get_arr(), max_obj)
return imgs, annos
def augmentation_func(
self, image: np.ndarray, target: Union[np.ndarray, int]
) -> Tuple[np.ndarray, Union[np.ndarray, int]]:
"""
Augmentation function that transforms image and target if augment_target is True
:param image:
:param target:
:return: returns a tuple of augmented images and target
"""
if self.augment_target:
# Check if target width and height make sense
assert image.shape[:-1] == target.shape, (
"Mask and image should be of the same "
"shape to be augmented! Mask is of shape {}".format(
target.shape))
segmap = SegmentationMapsOnImage(target, shape=image.shape[:-1])
image, target = self.sequence(image=image,
segmentation_maps=segmap)
target = target.get_arr()
else:
image = self.sequence(image=image)
return image, target
def augment_image(img, seg):
img = img.numpy().reshape(img.numpy().shape[-3:])
seg = seg.numpy().reshape(seg.numpy().shape[-3:])
shape = img.shape
segg = SegmentationMapsOnImage(seg, shape=shape)
a, b = self.seq.augment(image=img, segmentation_maps=segg)
return a, b.get_arr()
def segmaps_flipped(self):
segmaps_arr = np.int32([
[2, 2, 2],
[0, 1, 2],
[0, 1, 2],
])
return SegmentationMapsOnImage(segmaps_arr, shape=(3, 3, 3))
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 get_data(self, idx):
img = cv2.imread(self.data[idx][0])
anns = self.data[idx][1]
seg = np.zeros([img.shape[0], img.shape[1]], dtype=np.uint8)
for ann in anns:
points = np.int64(ann['segmentation']).reshape(-1, 2)
seg = cv2.fillPoly(seg, [points], ann['category_id'] + 1, 0)
seg = SegmentationMapsOnImage(seg, shape=img.shape)
return img, seg
def get_data(self, idx):
img = cv2.imread(self.data[idx][0])
seg_color = cv2.imread(self.data[idx][1])
seg = np.zeros([seg_color.shape[0], seg_color.shape[1]],
dtype=np.uint8)
for ci, c in enumerate(self.colormap):
seg[(seg_color == c).all(2)] = ci
seg = SegmentationMapsOnImage(seg, shape=img.shape)
return img, seg
def morph_close_mask_augmenter(segmaps, random_state, parents, hooks):
"""Morphological close augmenter"""
result = []
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
for segmap in segmaps:
arr = cv2.morphologyEx(segmap.arr.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
result.append(SegmentationMapsOnImage(arr, shape=arr.shape))
return result
def __getitem__(self, idx):
# images/101.png
image_name = self.image_list[idx]
image_path = os.path.join(self._base_dir, image_name)
image_obj = Image.open(image_path, 'r')
image = np.array(image_obj)
if image.ndim == 2:
image = np.expand_dims(image, 2).repeat(3, 2)
if image.shape[2] == 1:
image = image.repeat(3, 2)
if self.has_mask:
# masks/101.png
mask_name = image_name.replace("images", "masks")
mask_path = os.path.join(self._base_dir, mask_name)
mask = load_mask(mask_path, self.binarize)
else:
# Fake mask with all zeros. 'class' is always at the last dimension.
mask = np.zeros((1, ) + image.shape[:2], dtype=np.uint8)
mask_path = ""
unscaled_size = torch.tensor(image.shape[:2])
if self.mode == 'train' and self.train_loc_prob > 0 \
and np.random.random() < self.train_loc_prob:
image, mask = localize(image, mask, self.min_output_size)
if self.common_aug_func:
# Always mask == segmap.get_arr() == segmap.arr.astype(np.uint8).
segmap = SegmentationMapsOnImage(mask, shape=image.shape)
image_aug, segmap_aug = self.common_aug_func(
image=image, segmentation_maps=segmap)
image = image_aug
mask = segmap_aug.get_arr()
mask = mask.astype(np.uint8)
if self.image_trans_func:
image_obj = Image.fromarray(image)
image = self.image_trans_func(image_obj)
if self.segmap_trans_func:
mask = self.segmap_trans_func(mask)
else:
mask = mask.transpose([2, 0, 1])
sample = {
'image': image,
'mask': mask,
'image_path': image_path,
'mask_path': mask_path,
'crop_pos': torch.tensor((-1, -1)),
'unscaled_size': unscaled_size,
'uncropped_size': torch.tensor((-1, -1)),
'weight': self.ds_weight
}
return sample
def test_augment_segmaps__kernel_size_is_two__keep_size(self):
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
arr = np.int32([[0, 1, 2], [1, 2, 3]])
segmaps = SegmentationMapsOnImage(arr, shape=(6, 6, 3))
aug = self.augmenter(2, keep_size=True)
segmaps_aug = aug.augment_segmentation_maps(segmaps)
assert segmaps_aug.shape == (6, 6, 3)
assert np.allclose(segmaps_aug.arr, arr[..., np.newaxis])
def seg_augmentation(imgs, labs, crop_h=512, crop_w=1024):
aug_imgs, aug_labs = [], []
for img, lab in zip(imgs, labs):
img, lab = seg_seq_augmenter(crop_h, crop_w)(
image=img,
segmentation_maps=SegmentationMapsOnImage(lab, shape=img.shape))
lab = lab.get_arr().astype(np.uint8)
aug_imgs.append(img)
aug_labs.append(lab)
return np.array(aug_imgs), np.array(aug_labs)
def _segmentation_to_numpy(logits_or_mask, image=None):
"""
convert segmentation mask into numpy array, colors are taken from cityscapes.
:param logits_or_mask: tensor of shape C x H x W (logits) or 1 x H x W or H x W(mask)
:param image: None or Tensor of shape 3 x H x W. image to be used as background.
:return numpy array: shape 3 x H x W
"""
assert len(logits_or_mask.shape) == 3 or len(
logits_or_mask.shape) == 2, 'wrong input shape!'
if len(logits_or_mask.shape) == 2:
logits_or_mask = logits_or_mask.unsqueeze(0)
if logits_or_mask.shape[0] > 1:
mask = torch.argmax(logits_or_mask, dim=0,
keepdim=True).cpu().detach().numpy()
else:
mask = logits_or_mask.cpu().detach().numpy()
mask = mask.transpose((1, 2, 0))
# mask[mask==255] = 40 # dirty hack to move background class to 40
image_shape = (*mask.shape[:2], 3)
if image is not None:
image = image.cpu().detach().numpy().transpose((1, 2, 0))
assert image_shape == image.shape, 'Image shape and mask shape doesn\'t agree!'
else:
image = np.zeros(image_shape)
segmap = SegmentationMapsOnImage(mask.astype(np.int32),
shape=image_shape)
# if image is not None:
# out_image = segmap.draw_on_image(image)[0]
# else:
# out_image = segmap.draw(size=image_shape[:2])[0]
out_image = segmap.draw_on_image(
image.astype(np.uint8),
draw_background=False,
background_class_id=255,
alpha=1,
colors=Cityscapes.get_colors(use_train_id=True))[0]
out_image = out_image.transpose((2, 0, 1))
return out_image
def rot_2D_random(img, lbl):
segmap = SegmentationMapsOnImage(lbl.astype(np.int8, copy=False),
shape=lbl.shape)
# random_rotate = np.random.randint(90,size=1)
aug = iaa.Affine(rotate=(-90, 90))
img, lbl_aug = aug(image=img, segmentation_maps=segmap)
lbl = lbl_aug.get_arr()
return img, lbl
def apply_to_image_and_mask(augs, img, mask):
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
segmaps = SegmentationMapsOnImage(mask, shape=img.shape[:2])
res_img, _, res_segmaps = _apply(augs, img, masks=segmaps)
res_mask = res_segmaps.get_arr()
if res_img.shape[:2] != res_mask.shape[:2]:
raise ValueError(
f"Image and mask have different shapes "
f"({res_img.shape[:2]} != {res_mask.shape[:2]}) after augmentations. "
f"Please, contact tech support")
return res_img, res_mask
