|
from typing import Optional |
|
|
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
from matplotlib.pyplot import Figure |
|
|
|
|
|
def render_masks( |
|
image, |
|
masks, |
|
random_color: Optional[bool] = True, |
|
smoothen_contours: Optional[bool] = True, |
|
) -> "Figure": |
|
h, w = image.shape[:2] |
|
fig, ax = plt.subplots(figsize=(w / 100, h / 100), dpi=100) |
|
ax.axis("off") |
|
ax.imshow(image) |
|
|
|
for mask in masks: |
|
if random_color: |
|
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) |
|
else: |
|
color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6]) |
|
|
|
mask = mask.astype(np.uint8) |
|
mask = mask.reshape(h, w) |
|
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) |
|
|
|
if smoothen_contours: |
|
import cv2 |
|
|
|
contours, _ = cv2.findContours( |
|
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE |
|
) |
|
contours = [ |
|
cv2.approxPolyDP(contour, epsilon=0.01, closed=True) |
|
for contour in contours |
|
] |
|
mask_image = cv2.drawContours( |
|
mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2 |
|
) |
|
|
|
ax.imshow(mask_image, alpha=0.6) |
|
|
|
|
|
ax.set_xlim(0, w) |
|
ax.set_ylim(h, 0) |
|
plt.subplots_adjust(left=0, right=1, top=1, bottom=0) |
|
|
|
return fig |
|
|
