ROS wrapper for Meta's Segment-Anything model
#!/usr/bin/env python
import numpy as np
import rospy
import cv2

from cv_bridge   import CvBridge

from ros_sam.srv import Segmentation         as SegmentationSrv, \
                        SegmentationRequest  as SegmentationRequestMsg, \
                        SegmentationResponse as SegmentationResponseMsg

from ros_sam import SAM


if __name__ == '__main__':
    rospy.init_node('ros_sam')

    model = rospy.get_param('~model', 'vit_l')
    cuda  = rospy.get_param('~cuda', 'cuda')

    bridge = CvBridge()

    print('Starting SAM...')

    sam = SAM(model, cuda)

    def srv_segmentation(req : SegmentationRequestMsg):
        try:
            img    = cv2.cvtColor(bridge.imgmsg_to_cv2(req.image), cv2.COLOR_BGR2RGB)
            points = np.vstack([(p.x, p.y) for p in req.query_points])
            boxes  = np.asarray(req.boxes.data).reshape((len(req.boxes.data) // 4, 4))[0] if len(req.boxes.data) > 0 else None
            labels = np.asarray(req.query_labels)

            print("Segmenting at pixels:")
            for point in points:
                print(f"{point[0]}, {point[1]}")
                

            print(img.shape)

            masks, scores, logits = sam.segment(img, points, labels, boxes, req.multimask)

            res = SegmentationResponseMsg()
            res.masks  = [bridge.cv2_to_imgmsg(m.astype(np.uint8)) for m in masks]
            res.scores = scores.tolist()
            if req.logits:
                res.logits = [bridge.cv2_to_imgmsg(l) for l in logits]
            return res
        except Exception as e:
            print(f'{e}')
            raise Exception('Failure during service call. Check output on SAM node.')

    srv = rospy.Service('~segment', SegmentationSrv, srv_segmentation)

    print('SAM is ready')

    while not rospy.is_shutdown():
        rospy.sleep(0.2)
    
scripts/sam_node.py
robot-learning-freiburg/ros_sam
#!/usr/bin/env python
import cv2
import rospy
import numpy as np
import matplotlib.pyplot as plt

from cv_bridge import CvBridge
from pathlib   import Path

from geometry_msgs.msg import Point as PointMsg
from std_msgs.msg      import Int32MultiArray as Int32MultiArrayMsg
from ros_sam.srv import Segmentation         as SegmentationSrv, \
                        SegmentationRequest  as SegmentationRequestMsg


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))   


if __name__ == '__main__':
    rospy.init_node('ros_sam_test')

    print('Waiting for SAM service...')
    rospy.wait_for_service('ros_sam/segment')
    print('Found SAM service')

    image  = cv2.imread(f'{Path(__file__).parent}/../data/car.jpg')
    # image  = cv2.imread(f'{Path(__file__).parent}/../data/lab.png')
    image  = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    bridge  = CvBridge()
    img_msg = bridge.cv2_to_imgmsg(image)
    boxes = Int32MultiArrayMsg()
    
    points  = np.array([[566, 414]])

    msg_points  = [PointMsg(x=x, y=y, z=0) for (x, y) in points]
    labels      = np.array([1])

    sam_service = rospy.ServiceProxy('ros_sam/segment', SegmentationSrv)

    print(img_msg.header)
    print(img_msg.width)
    print(img_msg.height)
    print(img_msg.encoding)

    try:
        res = sam_service(img_msg, msg_points, labels, 
                          Int32MultiArrayMsg(), True, True)
    except rospy.ServiceException as e:
        print(f'{e}')
        print('exception raised')
        exit(0)

    for i, (msg_mask, score) in enumerate(zip(res.masks, res.scores)):
        mask = bridge.imgmsg_to_cv2(msg_mask)
        plt.figure(figsize=(10,10))
        plt.imshow(image)
        show_mask(mask, plt.gca())
        show_points(points, labels, plt.gca())
        plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
        plt.axis('off')
        plt.show()

scripts/sam_test.py
robot-learning-freiburg/ros_sam
#!/usr/bin/env python

# Try catkin install
from catkin_pkg.python_setup import generate_distutils_setup
from distutils.core import setup

from pathlib import Path


with open(f'{Path(__file__).parent}/requirements.txt', 'r') as f:
   pip_dependencies = f.readlines()


d = generate_distutils_setup(
   packages=['ros_sam'],
   package_dir={'': 'src'}
)

d.update({'install_requires': pip_dependencies})

setup(**d)

setup.py
robot-learning-freiburg/ros_sam
from .sam        import SAM
from .sam_client import SAMClient
from .utils      import show_box,  \
                        show_mask, \
                        show_points

src/ros_sam/__init__.py
robot-learning-freiburg/ros_sam
from pathlib import Path
import torch

from segment_anything import sam_model_registry, \
                             SamPredictor

# Free up GPU memory before loading the model
import gc


class SAM():
    def __init__(self, model_type, cuda_device=None) -> None:
        checkpoints = list(Path(f'{Path(__file__).parent}/../../models').glob(f'sam_{model_type}_*.pth'))

        if len(checkpoints) == 0:
            raise RuntimeError(f'No matching checkpoints for SAM model "{model_type}" was found in "package://ros_sam/models"')
        
        if len(checkpoints) > 1:
            raise RuntimeError(f'No unique checkpoint for SAM model "{model_type}" found in "package://ros_sam/models"')
        
        gc.collect()
        torch.cuda.empty_cache()

        sam_model    = sam_model_registry[model_type](checkpoint=checkpoints[0])
        self._device = cuda_device

        if self._device is not None:
            sam_model.to(device=self._device)

        self._predictor = SamPredictor(sam_model)
    
    def __del__(self):
        gc.collect()
        torch.cuda.empty_cache()

    def segment(self, img, points, point_labels, boxes=None, multimask=True):
        self._predictor.set_image(img)        
        return self._predictor.predict(
            point_coords=points,
            point_labels=point_labels,
            box=boxes,
            multimask_output=multimask,
        )

src/ros_sam/sam.py
robot-learning-freiburg/ros_sam
import rospy
import numpy as np

from cv_bridge import CvBridge

from geometry_msgs.msg import Point           as PointMsg
from std_msgs.msg      import Int32MultiArray as Int32MultiArrayMsg

from ros_sam.srv import Segmentation         as SegmentationSrv, \
                        SegmentationRequest  as SegmentationRequestMsg


class SAMClient():
    def __init__(self, service) -> None:
        self._bridge = CvBridge()
        
        rospy.wait_for_service(f'{service}/segment')

        self._srv_sam = rospy.ServiceProxy(f'{service}/segment', SegmentationSrv)

    def segment(self, img_rgb, points, labels, boxes=None):
        msg_boxes = Int32MultiArrayMsg()
        if boxes is not None:
            msg_boxes.data = list(boxes.flatten().astype(int))

        res = self._srv_sam(self._bridge.cv2_to_imgmsg(img_rgb), 
                            [PointMsg(x=x, y=y, z=0) for (x, y) in points],
                            labels, 
                            msg_boxes, 
                            boxes is None, False)
        return [self._bridge.imgmsg_to_cv2(m) for m in res.masks], res.scores

src/ros_sam/sam_client.py
robot-learning-freiburg/ros_sam
import numpy as np
import matplotlib.pyplot as plt

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))   
src/ros_sam/utils.py
robot-learning-freiburg/ros_sam