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)