mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
# This file contains utilities for Grounded-SAM mask refinement based on:
|
|
# https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
|
|
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
|
|
def mask_to_polygon(mask: npt.NDArray[np.uint8]) -> list[tuple[int, int]]:
|
|
"""Convert a binary mask to a polygon.
|
|
|
|
Returns:
|
|
list[list[int]]: List of (x, y) coordinates representing the vertices of the polygon.
|
|
"""
|
|
# Find contours in the binary mask.
|
|
contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
# Find the contour with the largest area.
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|
|
|
# Extract the vertices of the contour.
|
|
polygon = largest_contour.reshape(-1, 2).tolist()
|
|
|
|
return polygon
|
|
|
|
|
|
def polygon_to_mask(
|
|
polygon: list[tuple[int, int]], image_shape: tuple[int, int], fill_value: int = 1
|
|
) -> npt.NDArray[np.uint8]:
|
|
"""Convert a polygon to a segmentation mask.
|
|
|
|
Args:
|
|
polygon (list): List of (x, y) coordinates representing the vertices of the polygon.
|
|
image_shape (tuple): Shape of the image (height, width) for the mask.
|
|
fill_value (int): Value to fill the polygon with.
|
|
|
|
Returns:
|
|
np.ndarray: Segmentation mask with the polygon filled (with value 255).
|
|
"""
|
|
# Create an empty mask.
|
|
mask = np.zeros(image_shape, dtype=np.uint8)
|
|
|
|
# Convert polygon to an array of points.
|
|
pts = np.array(polygon, dtype=np.int32)
|
|
|
|
# Fill the polygon with white color (255).
|
|
cv2.fillPoly(mask, [pts], color=(fill_value,))
|
|
|
|
return mask
|