File size: 3,876 Bytes
0ca2a11 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
#!/usr/bin/env python
# coding: utf-8
###overlay
import cv2
import math
import random
import colorsys
import numpy as np
import itertools
import matplotlib.pyplot as plt
from matplotlib import cm
import os
import scipy.io as io
def get_bounding_box(img):
"""Get bounding box coordinate information."""
rows = np.any(img, axis=1)
cols = np.any(img, axis=0)
rmin, rmax = np.where(rows)[0][[0, -1]]
cmin, cmax = np.where(cols)[0][[0, -1]]
# due to python indexing, need to add 1 to max
# else accessing will be 1px in the box, not out
rmax += 1
cmax += 1
return [rmin, rmax, cmin, cmax]
####
def colorize(ch, vmin, vmax):
"""Will clamp value value outside the provided range to vmax and vmin."""
cmap = plt.get_cmap("jet")
ch = np.squeeze(ch.astype("float32"))
vmin = vmin if vmin is not None else ch.min()
vmax = vmax if vmax is not None else ch.max()
ch[ch > vmax] = vmax # clamp value
ch[ch < vmin] = vmin
ch = (ch - vmin) / (vmax - vmin + 1.0e-16)
# take RGB from RGBA heat map
ch_cmap = (cmap(ch)[..., :3] * 255).astype("uint8")
return ch_cmap
####
def random_colors(N, bright=True):
"""Generate random colors.
To get visually distinct colors, generate them in HSV space then
convert to RGB.
"""
brightness = 1.0 if bright else 0.7
hsv = [(i / N, 1, brightness) for i in range(N)]
colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
random.shuffle(colors)
return colors
####
def visualize_instances_map(
input_image, inst_map, type_map=None, type_colour=None, line_thickness=2
):
"""Overlays segmentation results on image as contours.
Args:
input_image: input image
inst_map: instance mask with unique value for every object
type_map: type mask with unique value for every class
type_colour: a dict of {type : colour} , `type` is from 0-N
and `colour` is a tuple of (R, G, B)
line_thickness: line thickness of contours
Returns:
overlay: output image with segmentation overlay as contours
"""
overlay = np.copy((input_image).astype(np.uint8))
inst_list = list(np.unique(inst_map)) # get list of instances
inst_list.remove(0) # remove background
inst_rng_colors = random_colors(len(inst_list))
inst_rng_colors = np.array(inst_rng_colors) * 255
inst_rng_colors = inst_rng_colors.astype(np.uint8)
for inst_idx, inst_id in enumerate(inst_list):
inst_map_mask = np.array(inst_map == inst_id, np.uint8) # get single object
y1, y2, x1, x2 = get_bounding_box(inst_map_mask)
y1 = y1 - 2 if y1 - 2 >= 0 else y1
x1 = x1 - 2 if x1 - 2 >= 0 else x1
x2 = x2 + 2 if x2 + 2 <= inst_map.shape[1] - 1 else x2
y2 = y2 + 2 if y2 + 2 <= inst_map.shape[0] - 1 else y2
inst_map_crop = inst_map_mask[y1:y2, x1:x2]
contours_crop = cv2.findContours(
inst_map_crop, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
# only has 1 instance per map, no need to check #contour detected by opencv
#print(contours_crop)
contours_crop = np.squeeze(
contours_crop[0][0].astype("int32")
) # * opencv protocol format may break
if len(contours_crop.shape) == 1:
contours_crop = contours_crop.reshape(1,-1)
#print(contours_crop.shape)
contours_crop += np.asarray([[x1, y1]]) # index correction
if type_map is not None:
type_map_crop = type_map[y1:y2, x1:x2]
type_id = np.unique(type_map_crop).max() # non-zero
inst_colour = type_colour[type_id]
else:
inst_colour = (inst_rng_colors[inst_idx]).tolist()
cv2.drawContours(overlay, [contours_crop], -1, inst_colour, line_thickness)
return overlay
# In[ ]:
|