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 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])
