Open sulaiman1988 opened 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
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.
I just want to know is it possible to find the contour inside ONNX model ?
How can I return just segmentation polygon points for each object?