segment-anything icon indicating copy to clipboard operation
segment-anything copied to clipboard

segmentation polygon

Open sulaiman1988 opened this issue 1 year ago • 3 comments

How can I return just segmentation polygon points for each object?

sulaiman1988 avatar Apr 09 '23 00:04 sulaiman1988

import cv2
import numpy as np
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
from pycocotools import mask as mask_utils

sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>)

for ann in masks:
    m = ann['segmentation']
    if isinstance(m, np.ndarray) and m.dtype == bool:
        m = mask_utils.encode(np.asfortranarray(m))
    elif isinstance(m, dict) and 'counts' in m and 'size' in m:
        pass  # Already in RLE format
    else:
        print("Invalid segmentation format:", m)
        continue
    
    mask = mask_utils.decode(m)
    contours, hierarchy = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours = [np.squeeze(contour) for contour in contours] # Convert contours to the correct shape
    contours = [np.atleast_2d(contour) for contour in contours]

Contours should be an array of coordinate pairs now which you should be able to work with?

For example, generate the d attribute of an svg path for a given contour like so...

def contour_to_svg_path(contour):
    path_commands = ""

    for i, point in enumerate(contour):
        point = point[0]  # Add this line to fix the indexing issue
        if i == 0:
            path_commands += f"M {point[0]},{point[1]} "
        else:
            path_commands += f"L {point[0]},{point[1]} "
    path_commands += "Z"
    return path_commands

FullStackSimon avatar Apr 09 '23 01:04 FullStackSimon

import cv2
import numpy as np
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
from pycocotools import mask as mask_utils

sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>)

for ann in masks:
    m = ann['segmentation']
    if isinstance(m, np.ndarray) and m.dtype == bool:
        m = mask_utils.encode(np.asfortranarray(m))
    elif isinstance(m, dict) and 'counts' in m and 'size' in m:
        pass  # Already in RLE format
    else:
        print("Invalid segmentation format:", m)
        continue
    
    mask = mask_utils.decode(m)
    contours, hierarchy = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours = [np.squeeze(contour) for contour in contours] # Convert contours to the correct shape
    contours = [np.atleast_2d(contour) for contour in contours]

Contours should be an array of coordinate pairs now which you should be able to work with?

For example, generate the d attribute of an svg path for a given contour like so...

def contour_to_svg_path(contour):
    path_commands = ""

    for i, point in enumerate(contour):
        point = point[0]  # Add this line to fix the indexing issue
        if i == 0:
            path_commands += f"M {point[0]},{point[1]} "
        else:
            path_commands += f"L {point[0]},{point[1]} "
    path_commands += "Z"
    return path_commands

A problem occurs on point = point[0] , which will only visualize the first point in the contour.

mu-cai avatar May 02 '23 03:05 mu-cai