def create_visualization(image, y_pred, y_true):
"""Creates the qualitative results to visualize
"""
# format prediction and ground truth
y_pred = np.argmax(y_pred, axis=-1).astype(np.uint8)
y_true = np.argmax(y_true, axis=-1).astype(np.uint8)
# format raw image
# put values between 0 and 255 (before it was between 0 and 1)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
image = image * 255.
image = image.astype(np.uint8)
# generate segmaps
y_pred = SegmentationMapsOnImage(y_pred, shape=image.shape)
y_true = SegmentationMapsOnImage(y_true, shape=image.shape)
grid_image = ia.draw_grid([
image,
y_true.draw_on_image(image)[0],
y_pred.draw_on_image(image)[0]
],
cols=3)
return grid_image
def create_visualization(pred_dep, gt_dep, pred_seg, gt_seg, img):
pred_dep = cv2.cvtColor(pred_dep, cv2.COLOR_GRAY2BGR)
pred_dep = (pred_dep * 255).astype(np.uint8)
gt_dep = cv2.cvtColor(gt_dep, cv2.COLOR_GRAY2BGR)
gt_dep = (gt_dep * 255).astype(np.uint8)
# dep_heatmap = dep_heatmap.astype(np.float32)/255.
# dep_overlay = dep_overlay.astype(np.float32)/255.
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = (img * 255).astype(np.uint8)
# generate segmaps
pred_seg = np.argmax(pred_seg, axis=-1).astype(np.uint8)
gt_seg = np.argmax(gt_seg, axis=-1).astype(np.uint8)
pred_seg = SegmentationMapsOnImage(pred_seg, shape=img.shape)
gt_seg = SegmentationMapsOnImage(gt_seg, shape=img.shape)
grid_image = ia.draw_grid(
[
img,
pred_seg.draw_on_image(img)[0],
gt_seg.draw_on_image(img)[0], pred_dep, gt_dep
# dep_heatmap,
# dep_overlay,
],
cols=5)
return grid_image
def chapter_examples_segmentation_maps_bool_small():
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 mask (bool, 128x128).
# Here, we arbitrarily place a square on the image.
segmap = np.zeros((128, 128, 1), dtype=bool)
segmap[28:71, 35:85, 0] = True
segmap = SegmentationMapsOnImage(segmap, shape=image.shape)
# Draw three columns: (1) original image,
# (2) original image with mask on top, (3) only mask
cells = [
image,
segmap.draw_on_image(image)[0],
segmap.draw(size=image.shape[:2])[0]
]
# Convert cells to a grid image and save.
grid_image = ia.draw_grid(cells, cols=3)
# imageio.imwrite("example_segmaps_bool.jpg", grid_image)
save("examples_segmentation_maps",
"bool_small.jpg",
grid_image,
quality=90)
def create_visualization_rgb(pred_seg, gt_seg, img):
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = (img * 255).astype(np.uint8)
# generate segmaps
pred_seg = np.argmax(pred_seg, axis=-1).astype(np.uint8)
gt_seg = np.argmax(gt_seg, axis=-1).astype(np.uint8)
pred_seg = SegmentationMapsOnImage(pred_seg, shape=img.shape)
gt_seg = SegmentationMapsOnImage(gt_seg, shape=img.shape)
grid_image = ia.draw_grid([
img,
pred_seg.draw_on_image(img)[0],
gt_seg.draw_on_image(img)[0],
],
cols=3)
return grid_image
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 _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 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 chapter_examples_segmentation_maps_simple():
import imageio
import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
ia.seed(1)
# Load an example image (uint8, 128x128x3).
image = ia.quokka(size=(128, 128), extract="square")
# Define an example segmentation map (int32, 128x128).
# Here, we arbitrarily place some squares on the image.
# Class 0 is our intended 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
segmap = SegmentationMapsOnImage(segmap, shape=image.shape)
# Define our augmentation pipeline.
seq = iaa.Sequential(
[
iaa.Dropout([0.05, 0.2]), # drop 5% or 20% of all pixels
iaa.Sharpen((0.0, 1.0)), # sharpen the image
iaa.Affine(
rotate=(-45,
45)), # rotate by -45 to 45 degrees (affects segmaps)
iaa.ElasticTransformation(
alpha=50, sigma=5) # apply water effect (affects segmaps)
],
random_order=True)
# Augment images and segmaps.
images_aug = []
segmaps_aug = []
for _ in range(5):
images_aug_i, segmaps_aug_i = seq(image=image,
segmentation_maps=segmap)
images_aug.append(images_aug_i)
segmaps_aug.append(segmaps_aug_i)
# We want to generate an image containing the original input image and
# segmentation maps before/after augmentation. (Both multiple times for
# multiple augmentations.)
#
# The whole image is supposed to have five columns:
# (1) original image,
# (2) original image with segmap,
# (3) augmented image,
# (4) augmented segmap on augmented image,
# (5) augmented segmap on its own in.
#
# We now generate the cells of these columns.
#
# Note that draw_on_image() and draw() both return lists of drawn
# images. Assuming that the segmentation map array has shape (H,W,C),
# the list contains C items.
cells = []
for image_aug, segmap_aug in zip(images_aug, segmaps_aug):
cells.append(image) # column 1
cells.append(segmap.draw_on_image(image)[0]) # column 2
cells.append(image_aug) # column 3
cells.append(segmap_aug.draw_on_image(image_aug)[0]) # column 4
cells.append(segmap_aug.draw(size=image_aug.shape[:2])[0]) # column 5
# Convert cells to a grid image and save.
grid_image = ia.draw_grid(cells, cols=5)
# imageio.imwrite("example_segmaps.jpg", grid_image)
save("examples_segmentation_maps", "simple.jpg", grid_image, quality=90)
images_aug.append(images_aug_i)
segmaps_aug.append(segmaps_aug_i)
# We want to generate an image containing the original input image and
# segmentation maps before/after augmentation. (Both multiple times for
# multiple augmentations.)
#
# The whole image is supposed to have five columns:
# (1) original image,
# (2) original image with segmap,
# (3) augmented image,
# (4) augmented segmap on augmented image,
# (5) augmented segmap on its own in.
#
# We now generate the cells of these columns.
#
# Note that draw_on_image() and draw() both return lists of drawn
# images. Assuming that the segmentation map array has shape (H,W,C),
# the list contains C items.
cells = []
for image_aug, segmap_aug in zip(images_aug, segmaps_aug):
cells.append(image) # column 1
cells.append(segmap.draw_on_image(image)[0]) # column 2
cells.append(image_aug) # column 3
cells.append(segmap_aug.draw_on_image(image_aug)[0]) # column 4
cells.append(segmap_aug.draw(size=image_aug.shape[:2])[0]) # column 5
# Convert cells to a grid image and save.
grid_image = ia.draw_grid(cells, cols=5)
imageio.imwrite("example_segmaps.jpg", grid_image)
pred_mask = pred_mask[:, :, np.newaxis].astype(bool)
# raw_label = raw_label[:,: np.newaxis]
pred_img_mask = pred_img[:, :, np.newaxis]
shape = pred_img_mask.shape
pred_img_mask = pred_img_mask == 2
if np.sum(pred_img_mask) < 40:
pred_img_mask = np.zeros((512, 720, 1)).astype(np.bool)
# pred_mask += 1
# raw_label += 1
# aug = iaa.AddToBrightness(add=(30))
segmap = SegmentationMapsOnImage(pred_mask, shape=img.shape)
segmap = segmap.pad(top=0, right=0, bottom=0, left=0)
overlayed_seg = segmap.draw_on_image(img,
alpha=0.6,
colors=[(0, 0, 0), (0, 0, 255)])[0]
# overlayed_seg = aug(overlayed_seg)
# segmap_img = SegmentationMapsOnImage(pred_img_mask, shape=img.shape)
# overlayed_seg_img = segmap_img.draw_on_image(img, alpha=0.6, colors=[(0, 0, 0), (0, 0, 255)])[0]
# label_seg = SegmentationMapsOnImage(raw_label, shape=img.shape)
# overlayed_label = label_seg.draw_on_image(img,alpha=0.5,colors=[(0,0,0),(0,0,0),(0,0,153),(0,204,0)])[0]
heatmap = HeatmapsOnImage(confidence_map,
shape=img.shape,
min_value=0,
max_value=1)
new_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
overlayed_heatmap = heatmap.draw_on_image(new_img, alpha=0.45,
cmap='jet')[0]
overlayed_heatmap = cv2.cvtColor(overlayed_heatmap, cv2.COLOR_RGB2BGR)
