File size: 3,900 Bytes
d33e75e |
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 118 119 120 121 122 123 124 125 126 127 |
"""
Mask and point painting utilities
Adapted from MatAnyone demo
"""
import cv2
import numpy as np
from PIL import Image
def mask_painter(input_image, input_mask, mask_color=5, mask_alpha=0.7,
contour_color=1, contour_width=5):
"""
Paint mask on image with transparency
Args:
input_image: np.ndarray, (H, W, 3)
input_mask: np.ndarray, (H, W), binary mask
mask_color: int, color ID for mask
mask_alpha: float, transparency
contour_color: int, color ID for contour
contour_width: int, width of contour
Returns:
painted_image: np.ndarray, (H, W, 3)
"""
assert input_image.shape[:2] == input_mask.shape, "Image and mask must have same dimensions"
# Color palette
palette = np.array([
[0, 0, 0], # 0: black
[255, 0, 0], # 1: red
[0, 255, 0], # 2: green
[0, 0, 255], # 3: blue
[255, 255, 0], # 4: yellow
[255, 0, 255], # 5: magenta
[0, 255, 255], # 6: cyan
[128, 128, 128], # 7: gray
[255, 165, 0], # 8: orange
[128, 0, 128], # 9: purple
])
mask_color_rgb = palette[mask_color % len(palette)]
contour_color_rgb = palette[contour_color % len(palette)]
# Create colored mask
painted_image = input_image.copy()
colored_mask = np.zeros_like(input_image)
colored_mask[input_mask > 0] = mask_color_rgb
# Blend with alpha
mask_region = input_mask > 0
painted_image[mask_region] = (
painted_image[mask_region] * (1 - mask_alpha) +
colored_mask[mask_region] * mask_alpha
).astype(np.uint8)
# Draw contour
if contour_width > 0:
contours, _ = cv2.findContours(
input_mask.astype(np.uint8),
cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE
)
cv2.drawContours(
painted_image,
contours,
-1,
contour_color_rgb.tolist(),
contour_width
)
return painted_image
def point_painter(input_image, input_points, point_color=8, point_alpha=0.9,
point_radius=15, contour_color=2, contour_width=3):
"""
Paint points on image
Args:
input_image: np.ndarray, (H, W, 3)
input_points: np.ndarray, (N, 2), [x, y] coordinates
point_color: int, color ID for points
point_alpha: float, transparency
point_radius: int, radius of point circles
contour_color: int, color ID for contour
contour_width: int, width of contour
Returns:
painted_image: np.ndarray, (H, W, 3)
"""
if len(input_points) == 0:
return input_image
palette = np.array([
[0, 0, 0], # 0: black
[255, 0, 0], # 1: red
[0, 255, 0], # 2: green
[0, 0, 255], # 3: blue
[255, 255, 0], # 4: yellow
[255, 0, 255], # 5: magenta
[0, 255, 255], # 6: cyan
[128, 128, 128], # 7: gray
[255, 165, 0], # 8: orange
[128, 0, 128], # 9: purple
])
point_color_rgb = palette[point_color % len(palette)]
contour_color_rgb = palette[contour_color % len(palette)]
painted_image = input_image.copy()
for point in input_points:
x, y = int(point[0]), int(point[1])
# Draw filled circle with alpha blending
overlay = painted_image.copy()
cv2.circle(overlay, (x, y), point_radius, point_color_rgb.tolist(), -1)
cv2.addWeighted(overlay, point_alpha, painted_image, 1 - point_alpha, 0, painted_image)
# Draw contour
if contour_width > 0:
cv2.circle(painted_image, (x, y), point_radius, contour_color_rgb.tolist(), contour_width)
return painted_image
|