Spaces:
Running
Running
File size: 2,116 Bytes
650c5f6 |
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 |
import numpy as np
import cv2
def overlay_predictions(img, mask=None, polygons=None, bbox=None):
overlayed = img.copy()
if bbox is not None:
overlayed = draw_bbox(overlayed, bbox)
if mask is not None:
overlayed = overlay_davis(overlayed, mask)
if polygons is not None:
overlayed = plot_polygons(overlayed, polygons)
return overlayed
def overlay_davis(image, mask, colors=[[0, 0, 0], [255, 102, 102]], cscale=1, alpha=0.4):
from scipy.ndimage.morphology import binary_dilation
colors = np.reshape(colors, (-1, 3))
colors = np.atleast_2d(colors) * cscale
im_overlay = image.copy()
object_ids = np.unique(mask)
h_i, w_i = image.shape[0:2]
h_m, w_m = mask.shape[0:2]
if h_i != h_m:
mask = cv2.resize(mask, [h_i, w_i], interpolation=cv2.INTER_NEAREST)
for object_id in object_ids[1:]:
# Overlay color on binary mask
foreground = image*alpha + np.ones(image.shape)*(1-alpha) * np.array(colors[object_id])
binary_mask = mask == object_id
# Compose image
im_overlay[binary_mask] = foreground[binary_mask]
# countours = skimage.morphology.binary.binary_dilation(binary_mask) - binary_mask
countours = binary_dilation(binary_mask) ^ binary_mask
# countours = cv2.dilate(binary_mask, cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))) - binary_mask
im_overlay[countours, :] = 0
return im_overlay.astype(image.dtype)
def draw_bbox(img, box, color=(0, 255, 0), thickness=3):
x1, y1, x2, y2 = box
return cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness=thickness)
def plot_polygons(img, polygons, radius=3):
for polygon in polygons:
if len(polygon) > 0:
polygon = np.reshape(polygon[:len(polygon)-len(polygon)%2], (len(polygon)//2, 2)).astype(np.int16)
for i, point in enumerate(polygon):
color = (255, 0, 0)
#if i == 0:
# color = (0, 0, 255)
img = cv2.circle(img, point, radius, color, thickness=-1)
return img |