diff --git a/README.md b/README.md index 9a186b2..14b03c1 100644 --- a/README.md +++ b/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**
> 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). diff --git a/docs/instance_seg.md b/docs/instance_seg.md index e9fb806..faeed0d 100644 --- a/docs/instance_seg.md +++ b/docs/instance_seg.md @@ -25,7 +25,7 @@ source /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 diff --git a/semantic_inference/python/semantic_inference/models/instance_segmenter.py b/semantic_inference/python/semantic_inference/models/instance_segmenter.py index 946ef7b..7b8137e 100644 --- a/semantic_inference/python/semantic_inference/models/instance_segmenter.py +++ b/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, ) diff --git a/semantic_inference/python/semantic_inference/models/wrappers.py b/semantic_inference/python/semantic_inference/models/wrappers.py index bdd6206..2a4c237 100644 --- a/semantic_inference/python/semantic_inference/models/wrappers.py +++ b/semantic_inference/python/semantic_inference/models/wrappers.py @@ -346,23 +346,25 @@ 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 @@ -370,42 +372,44 @@ 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): @@ -413,6 +417,61 @@ def load(cls, filepath): 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( diff --git a/semantic_inference/python/semantic_inference/visualization.py b/semantic_inference/python/semantic_inference/visualization.py index fa193d4..44cc7ed 100644 --- a/semantic_inference/python/semantic_inference/visualization.py +++ b/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) diff --git a/semantic_inference_ros/app/instance_segmentation_node b/semantic_inference_ros/app/instance_segmentation_node index 5d4fba9..82de733 100755 --- a/semantic_inference_ros/app/instance_segmentation_node +++ b/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() diff --git a/semantic_inference_ros/config/instance_segmentation/gdsam2.yaml b/semantic_inference_ros/config/instance_segmentation/gdsam2.yaml index f943275..6b08789 100644 --- a/semantic_inference_ros/config/instance_segmentation/gdsam2.yaml +++ b/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 diff --git a/semantic_inference_ros/config/instance_segmentation/yolov11.yaml b/semantic_inference_ros/config/instance_segmentation/yolo-seg.yaml similarity index 82% rename from semantic_inference_ros/config/instance_segmentation/yolov11.yaml rename to semantic_inference_ros/config/instance_segmentation/yolo-seg.yaml index cef82eb..c2852e1 100644 --- a/semantic_inference_ros/config/instance_segmentation/yolov11.yaml +++ b/semantic_inference_ros/config/instance_segmentation/yolo-seg.yaml @@ -1,6 +1,6 @@ --- model: instance_model: - type: yolov11 + type: yolo-seg model_name: yolo11n-seg.pt confidence_threshold: 0.25 diff --git a/semantic_inference_ros/config/instance_segmentation/yoloe.yaml b/semantic_inference_ros/config/instance_segmentation/yoloe.yaml new file mode 100644 index 0000000..902b3d8 --- /dev/null +++ b/semantic_inference_ros/config/instance_segmentation/yoloe.yaml @@ -0,0 +1,8 @@ +--- +model: + instance_model: + type: yoloe + model_name: yoloe-26m-seg.pt + text_prompt: [ignore, chair, cone, bag] + confidence_threshold: 0.1 + min_segmentation_size: 0