Feature/bump hydra instance seg

README.md

@@ -80,6 +80,7 @@ echo "build: {cmake-args: [--no-warn-unused-cli, -DCMAKE_BUILD_TYPE=Release]}" >
 Once you've added this repository to your workspace, follow one (or both) of the following setup-guides as necessary:
 - [Closed-Set](docs/closed_set.md#setting-up)
 - [Open-Set](docs/open_set.md#setting-up)
+- [Instance-Seg](docs/instance_seg.md)
 
 > **Note** </br>
 > Some of our other (larger) packages have or will have more accessible guides to getting `semantic_inference` set up for specific applications, such as [Hydra](https://github.com/MIT-SPARK/Hydra), [Khronos](https://github.com/MIT-SPARK/Khronos) or [Clio](https://github.com/MIT-SPARK/Clio).

docs/instance_seg.md

@@ -25,7 +25,7 @@ source <PATH_TO_ENVIRONMENT>/bin/activate
 pip install ./semantic_inference
 ```
 
-The above setup allows you to use `yolov11`, in order to use `Grounded Sam 2`, we have to manually install it.
+The above setup allows you to use `yolo instance segmentation`, in order to use `Grounded Sam 2`, we have to manually install it.
 ```shell
 # cd to your favorite path, we can default to `~/.semantic_inference/`
 git clone -b more_gpu https://github.com/MultyXu/Grounded-SAM-2.git
@@ -35,7 +35,9 @@ And follow the `README.md` in the cloned repo to install gdsam2.
 ## Setup models
 For `Grounded Sam 2`, put (or symlink) `GroundingDINO_SwinT_OGC.py` under `~/.semantic_inference/gdsam2_config/`. And, put `sam2.1_hiera_large.pt` and `groundingdino_swint_ogc.pth` under `~/.semantic_inference/`
 
-For `Yolo`, download [yolo11n-seg.pt](https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-seg.pt) under `~/.semantic_inference/`
+For `YOLO`, download [yolo11n-seg.pt](https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n-seg.pt) under `~/.semantic_inference/`.
+
+For `YOLOE`, download [yoloe-26m-seg.pt](https://github.com/ultralytics/assets/releases/download/v8.4.0/yoloe-26m-seg.pt) under `~/.semantic_inference/`.
 
 ## Trying out close-set instance segmentation nodes
 

semantic_inference/python/semantic_inference/models/instance_segmenter.py

@@ -52,6 +52,7 @@ class Results:
     boxes: torch.Tensor  # (n, 4) xyxy format, torch.float32
     categories: torch.Tensor  # (n,), torch.float32/int64 (doesn't matter)
     confidences: torch.Tensor  # (n,), torch.float32
+    img_shape: tuple[int, int]
 
     @property
     def instance_seg_img(self):
@@ -60,6 +61,9 @@ def instance_seg_img(self):
         Each pixel value encodes both category id and instance id.
         First 16 bits are category id, last 16 bits are instance id.
         """
+        if self.masks is None:
+            return np.zeros(self.img_shape)
+
         masks = self.masks.cpu().numpy()
         category_ids = self.categories.cpu().numpy()
         img = np.zeros(masks[0].shape, dtype=np.uint32)
@@ -89,7 +93,7 @@ class InstanceSegmenterConfig(Config):
     """Main config for instance segmenter."""
 
     # relevant configs (model path, model weights) for the model
-    instance_model: Any = config_field("instance_model", default="yolov11")
+    instance_model: Any = config_field("instance_model", default="yolo-seg")
 
 
 class InstanceSegmenter(nn.Module):
@@ -152,8 +156,12 @@ def forward(self, rgb_img):
         Returns:
             Encoded image
         """
-        categories, masks, boxes, confidences = self.segmenter(rgb_img)
+        categories, masks, boxes, confidences, img_shape = self.segmenter(rgb_img)
 
         return Results(
-            masks=masks, boxes=boxes, categories=categories, confidences=confidences
+            masks=masks,
+            boxes=boxes,
+            categories=categories,
+            confidences=confidences,
+            img_shape=img_shape,
         )

semantic_inference/python/semantic_inference/models/wrappers.py

@@ -346,73 +346,132 @@ def load(cls, filepath):
         return Config.load(cls, filepath)
 
 
-class Yolov11InstanceSegmenterWrapper(nn.Module):
-    """Yolov11 instance segmentation wrapper."""
+class YoloInstanceSegmenterBase(nn.Module):
+    """Base class for YOLO instance segmentation wrappers.
+
+    Subclasses must assign self.model in their __init__.
+    """
 
     def __init__(self, config):
-        """Load Yolov11 model."""
+        """Store shared configuration."""
         super().__init__()
-        from ultralytics import YOLO
-
         self.config = config
-        model_weights = os.path.join(
-            path_to_dot_semantic_inference(), f"{config.model_name}"
-        )
-        self.model = YOLO(model_weights)
         self.confidence_threshold = config.confidence_threshold
+        self.min_segmentation_size = config.min_segmentation_size
+        self.overlap_merge_iou = config.overlap_merge_iou
+        self.model_weights = os.path.join(
+            path_to_dot_semantic_inference(), f"{self.config.model_name}"
+        )
 
     def eval(self):
-        """override eval to avoid issues with yolo model"""
+        """Override eval to avoid issues with yolo model."""
         self.model.model.eval()
 
     @property
     def category_names(self):
         """Get category names."""
         return self.model.names
 
-    @classmethod
-    def construct(cls, **kwargs):
-        """Load model from configuration dictionary."""
-        config = Yolov11InstanceSegmenterConfig()
-        config.update(kwargs)
-        return cls(config)
-
     def forward(self, img):
         """Segment image."""
-        result = self.model(img)[0]  # assume batch size 1
+        result = self.model(
+            img,
+            conf=self.confidence_threshold,
+            imgsz=(img.shape[0], img.shape[1]),
+            # Set True to maintain original image size for masks,
+            # required for hydra, but slow down the model
+            retina_masks=True,
+            iou=self.overlap_merge_iou,
+        )[0]  # assume batch size 1
         if result.masks is None:
-            return None, None, None, None
+            return None, None, None, None, (img.shape[0], img.shape[1])
 
-        categories = result.boxes.cls.cpu()  # int8
+        categories = result.boxes.cls.to(torch.int).cpu()  # int8
         masks = result.masks.data.to(torch.bool).cpu()
         boxes = result.boxes.xyxy.cpu()  # float32
         confidences = result.boxes.conf.cpu()  # float32
 
-        # filter by confidence threshold
-        mask_indices = confidences > self.confidence_threshold
-        categories = categories[mask_indices]
-        masks = masks[mask_indices]
-        boxes = boxes[mask_indices]
-        confidences = confidences[mask_indices]
-        return categories, masks, boxes, confidences
+        # filter out small masks
+        if self.min_segmentation_size > 0:
+            mask_sizes = masks.sum(dim=(1, 2))  # number of pixels in each mask
+            size_indices = mask_sizes > self.min_segmentation_size
+            categories = categories[size_indices]
+            masks = masks[size_indices]
+            boxes = boxes[size_indices]
+            confidences = confidences[size_indices]
+        return categories, masks, boxes, confidences, (masks.shape[1], masks.shape[2])
 
 
-@register_config(
-    "instance_model", name="yolov11", constructor=Yolov11InstanceSegmenterWrapper
-)
 @dataclasses.dataclass
-class Yolov11InstanceSegmenterConfig(Config):
-    """Configuration for Yolov11 instance segmenter."""
+class YoloInstanceSegmenterBaseConfig(Config):
+    """Shared configuration for YOLO instance segmenters."""
 
-    model_name: str = "yolo11n-seg.pt"
+    model_name: str = ""
     confidence_threshold: float = 0.25
+    min_segmentation_size: int = 0  # number of pixels to filter out small masks
+    overlap_merge_iou: float = 0.7  # merging overlapping object detection
 
     @classmethod
     def load(cls, filepath):
         """Load config from file."""
         return Config.load(cls, filepath)
 
 
+class YolosegInstanceSegmenterWrapper(YoloInstanceSegmenterBase):
+    """Yolov11 instance segmentation wrapper."""
+
+    def __init__(self, config):
+        """Load Yolov11 model."""
+        from ultralytics import YOLO
+
+        super().__init__(config)
+        self.model = YOLO(self.model_weights)
+
+    @classmethod
+    def construct(cls, **kwargs):
+        """Load model from configuration dictionary."""
+        config = YolosegInstanceSegmenterConfig()
+        config.update(kwargs)
+        return cls(config)
+
+
+@register_config(
+    "instance_model", name="yolo-seg", constructor=YolosegInstanceSegmenterWrapper
+)
+@dataclasses.dataclass
+class YolosegInstanceSegmenterConfig(YoloInstanceSegmenterBaseConfig):
+    """Configuration for YOLO segmentation instance segmenter."""
+
+    model_name: str = "yolo11n-seg.pt"
+
+
+class YoloeInstanceSegmenterWrapper(YoloInstanceSegmenterBase):
+    """Yoloe instance segmentation wrapper."""
+
+    def __init__(self, config):
+        from ultralytics import YOLOE
+
+        if not config.text_prompt:
+            raise ValueError(
+                "text_prompt must be provided for YoloeInstanceSegmenterWrapper"
+            )
+
+        super().__init__(config)
+        self.model = YOLOE(self.model_weights)
+        self.model.set_classes(self.config.text_prompt)
+
+
+@register_config(
+    "instance_model", name="yoloe", constructor=YoloeInstanceSegmenterWrapper
+)
+@dataclasses.dataclass
+class YoloeInstanceSegmenterConfig(YoloInstanceSegmenterBaseConfig):
+    """Configuration for Yoloe instance segmenter."""
+
+    model_name: str = "yoloe-26m-seg.pt"
+    text_prompt: list[str] = dataclasses.field(default_factory=list)
+
+
 class GDSam2InstanceSegmenterWrapper(nn.Module):
     """Grounded SAM 2 instance segmentation wrapper."""
 
@@ -425,6 +484,12 @@ def __init__(self, config):
 
         self.config = config
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Enable TF32 for faster fp32 matrix ops on Ampere+ GPUs (Blackwell: major=10)
+        if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
         self.text_prompt = config.text_prompt
         self.multimask_output = config.multimask_output
         self.erosion = config.erosion
@@ -512,7 +577,7 @@ def forward(self, img):
 
         # if nothing detected
         if boxes.shape[0] == 0:
-            return None, None, None, None
+            return None, None, None, None, (img.shape[0], img.shape[1])
 
         # process the box prompt for SAM 2
         h, w, _ = img.shape
@@ -530,14 +595,15 @@ def forward(self, img):
             torch.backends.cudnn.allow_tf32 = True
         """
 
-        # SAM 2 predicts mask
-        self.sam2_predictor.set_image(img)
-        masks, scores, logits = self.sam2_predictor.predict(
-            point_coords=None,
-            point_labels=None,
-            box=input_boxes,
-            multimask_output=self.multimask_output,
-        )
+        # SAM2 predicts mask — bfloat16 autocast scoped to SAM2 only (leaves GD in fp32)
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
+            self.sam2_predictor.set_image(img)
+            masks, scores, logits = self.sam2_predictor.predict(
+                point_coords=None,
+                point_labels=None,
+                box=input_boxes,
+                multimask_output=self.multimask_output,
+            )
 
         # Sample best according to scores if multimask output
         if self.multimask_output:
@@ -575,7 +641,7 @@ def forward(self, img):
         # use xyxy boxes
         boxes = torch.tensor(input_boxes)
 
-        return categories, masks, boxes, confidences
+        return categories, masks, boxes, confidences, (masks.shape[1], masks.shape[2])
 
 
 @register_config(

semantic_inference/python/semantic_inference/visualization.py

@@ -328,6 +328,10 @@ def get_semantic_overlay_img(category_names, ret, img):
     boxes = ret.boxes
     confidences = ret.confidences
 
+    if masks is None:
+        # no bounding boxes to draw
+        return img
+
     # Convert RGB to BGR for OpenCV
     vis_img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 

semantic_inference_ros/app/instance_segmentation_node

@@ -1,7 +1,10 @@
 #!/usr/bin/env python3
 """Node that runs openset segmentation."""
 
+import collections
+import json
 import pathlib
+import time
 from dataclasses import dataclass, field
 
 import numpy as np
@@ -10,6 +13,7 @@ import spark_config as sc
 import torch
 from rclpy.node import Node
 from sensor_msgs.msg import Image
+from std_msgs.msg import String
 
 import semantic_inference.models as models
 import semantic_inference_ros
@@ -53,6 +57,11 @@ class InstanceSegmentationNode(Node):
         self._model.eval()  # TODO: causing issue with yolo model
         self.get_logger().info("Finished initializing!")
 
+        self._last_inference_time = None
+        self._inference_times_ms = collections.deque(maxlen=30)
+        self._status_pub = self.create_publisher(String, "~/status", 1)
+        self._hb_timer = self.create_timer(0.1, self._hb_callback)
+
         self._pub = self.create_publisher(Image, "semantic/image_raw", 1)
         self._worker = semantic_inference_ros.ImageWorker(
             self, self.config.worker, "color/image_raw", self._spin_once
@@ -64,18 +73,25 @@ class InstanceSegmentationNode(Node):
             )
 
     def _spin_once(self, header, img):
+        self.get_logger().debug(f"recived image of size: {img.shape}")
+        t0 = time.monotonic()
         with torch.no_grad():
-            ret = self._model.segment(img, is_rgb_order=True).cpu()
+            ret = self._model.segment(img, is_rgb_order=True)
+        self._inference_times_ms.append((time.monotonic() - t0) * 1e3)
 
         if ret.masks is None:
             self.get_logger().debug("No masks detected in the image.")
-            return
+        else:
+            self.get_logger().debug("get image size: " + str(ret.img_shape))
+            self.get_logger().debug(f"result category type: {ret.categories.dtype}")
 
+        # get 32 bit instance + cateogry img, all 0 if no box detected
         instance_seg_img = ret.instance_seg_img
         # Convert to int32 to match 32SC1 encoding expected by cv_bridge
         instance_seg_img = instance_seg_img.astype(np.int32)
         msg = Conversions.to_image_msg(header, instance_seg_img, encoding="32SC1")
         self._pub.publish(msg)
+        self._last_inference_time = self.get_clock().now().nanoseconds * 1e-9
         self.get_logger().debug("Published instance segmentation image.")
 
         if self.config.visualize_semantic_img:
@@ -85,6 +101,26 @@ class InstanceSegmentationNode(Node):
                 Conversions.to_image_msg(header, color_img, encoding="rgb8")
             )
 
+    def _hb_callback(self):
+        record = {
+            "nickname": "instance_seg",
+            "node_name": self.get_fully_qualified_name(),
+        }
+        now = self.get_clock().now().nanoseconds * 1e-9
+        if self._last_inference_time is None:
+            record["status"] = "STARTUP"
+            record["note"] = "Waiting for first image."
+        elif now - self._last_inference_time > 10.0:
+            record["status"] = "WARNING"
+            record["note"] = f"No inference for {now - self._last_inference_time:.1f}s."
+        else:
+            record["status"] = "NOMINAL"
+            avg_ms = sum(self._inference_times_ms) / len(self._inference_times_ms)
+            record["note"] = f"Average inference: {avg_ms:.1f} ms"
+        msg = String()
+        msg.data = json.dumps(record)
+        self._status_pub.publish(msg)
+
     def stop(self):
         """Stop the underlying image worker."""
         self._worker.stop()

semantic_inference_ros/config/instance_segmentation/gdsam2.yaml

@@ -6,8 +6,8 @@ model:
     text_prompt: unknown. green chair. van. fire hydrant. ground. grass. sand. sidewalk. dock. road. path. ignore. building. shelter. signal. rock. fence. boat. sign. hill. bridge. wall. floor. ceiling.
       door. stairs. pole. rail. structure. window. surface. flora. flower. bed. box. storage. barrel. bag. basket. ignore. flag. decor. light. appliance. trash. bicycle. food. clothes.
     # text_prompt: unknown. adirondack chair. van. fire hydrant.
-    sam2_checkpoint: sam2.1_hiera_large.pt
-    sam2_model_config: sam2.1_hiera_l.yaml
+    sam2_checkpoint: sam2.1_hiera_small.pt
+    sam2_model_config: sam2.1_hiera_s.yaml
     grounding_dino_config: GroundingDINO_SwinT_OGC.py
     grounding_dino_checkpoint: groundingdino_swint_ogc.pth
     box_threshold: 0.45