File size: 2,999 Bytes
8260e47 bf29adc 6038d30 8260e47 410698b bf29adc 410698b 6038d30 8260e47 410698b 6038d30 8260e47 6038d30 410698b 6038d30 410698b 6038d30 410698b 6038d30 410698b 6038d30 410698b 6038d30 8260e47 6038d30 8260e47 6038d30 410698b 6038d30 410698b 6038d30 |
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 |
from typing import Optional
import cv2
import numpy as np
from PIL import Image
def export_mask(
masks,
autogenerated: Optional[bool] = False,
random_color: Optional[bool] = True,
smoothen_contours: Optional[bool] = True,
) -> Image:
if not autogenerated:
num_masks, _, h, w = masks.shape
num_masks = len(masks)
# Ensure masks are 2D by squeezing channel dimension
masks = masks.squeeze(axis=1)
# Create a single uint8 image with unique values for each mask
combined_mask = np.zeros((h, w), dtype=np.uint8)
for i in range(num_masks):
mask = masks[i]
mask = mask.astype(np.uint8)
combined_mask[mask > 0] = i + 1
# Create color map for visualization
if random_color:
colors = np.random.rand(num_masks, 3) # Random colors for each mask
else:
colors = np.array(
[[30 / 255, 144 / 255, 255 / 255]] * num_masks
) # Use fixed color
# Create an RGB image where each mask has its own color
color_image = np.zeros((h, w, 3), dtype=np.uint8)
for i in range(1, num_masks + 1):
mask_color = colors[i - 1] * 255
color_image[combined_mask == i] = mask_color
# Convert the NumPy array to a PIL Image
pil_image = Image.fromarray(color_image)
# Optional: Add contours to the mask image
if smoothen_contours:
contours_image = np.zeros((h, w, 4), dtype=np.float32)
for i in range(1, num_masks + 1):
mask = (combined_mask == i).astype(np.uint8)
contours_image = smoothen(mask, contours_image)
# Convert contours to PIL image and blend with the color image
contours_image = (contours_image[:, :, :3] * 255).astype(np.uint8)
contours_pil_image = Image.fromarray(contours_image)
pil_image = Image.blend(pil_image, contours_pil_image, alpha=0.6)
return pil_image
else:
sorted_anns = sorted(masks, key=(lambda x: x["area"]), reverse=True)
img_shape = sorted_anns[0]["segmentation"].shape
img = np.ones((img_shape[0], img_shape[1], 4))
img[:, :, 3] = 0
for ann in sorted_anns:
m = ann["segmentation"]
color_mask = np.concatenate([np.random.random(3), [0.5]])
img[m] = color_mask
if smoothen_contours:
img = smoothen(m, img)
img = (img * 255).astype(np.uint8)
pil_image = Image.fromarray(img)
return pil_image
def smoothen(mask: np.ndarray, image: np.ndarray) -> np.ndarray:
contours, _ = cv2.findContours(
mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
)
contours = [
cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours
]
image = cv2.drawContours(image, contours, -1, (0, 0, 1, 0.4), thickness=1)
return image
|