facebookresearch / segment-anything

The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.
Apache License 2.0
47.39k stars 5.61k forks source link

segmentation polygon #121

Open sulaiman1988 opened 1 year ago

sulaiman1988 commented 1 year ago

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

FullStackSimon commented 1 year ago
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
mu-cai commented 1 year ago
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.

hariharan1412 commented 1 month ago

I just want to know is it possible to find the contour inside ONNX model ?