How to generate polygons from the mask ?

[mohamedusmb]

Hey guys, with this code, I want to generate some polygons to obtain coordinates and create a bounding box around the mask. Is this possible?

import numpy as np
import cv2
import torch
import matplotlib.pyplot as plt
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry

def segment_image(image, return_segmented_image=False):
    sam = sam_model_registry["vit_h"](checkpoint="checkpoints/sam_vit_h_4b8939.pth")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    sam.to(device=device)
    mask_generator = SamAutomaticMaskGenerator(model=sam)
   
    if image.shape[-1] == 3:  
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    masks = mask_generator.generate(image)
    if return_segmented_image:
        overlay_image = np.copy(image)
        for mask in masks:
            overlay_image[mask['segmentation']] = np.random.randint(0, 256, 3)   
        response = input("Voulez-vous afficher l'image segmentée ? (Oui/Non): ").lower()
        if response in ["oui"]:
             plt.imshow(overlay_image)
             plt.axis('off')
             plt.show()

    return masks

image_path = 'images/dog.jpg'
image = cv2.imread(image_path)
segment_image(image, return_segmented_image=True)

[Abonia1]

def mask_to_polygon(mask: np.ndarray) -> List[List[int]]:
    # 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]) -> np.ndarray:
    """
    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.

    Returns:
    - np.ndarray: Segmentation mask with the polygon filled.
    """
    # 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=(255,))

    return mask

Hope this may help.