Feature/labelspaces

install/download_weights.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+import argparse
+import pathlib
+
+import ultralytics
+
+import semantic_inference
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Downloading model weights")
+    parser.add_argument("models", nargs="*")
+    parser.add_argument("-d", "--directory", default=None)
+    args = parser.parse_args()
+
+    target_path = args.directory or semantic_inference.models_path()
+    target_path = pathlib.Path(target_path).expanduser().absolute()
+    target_path.mkdir(exist_ok=True, parents=True)
+    for model in args.models:
+        output_path = target_path / f"{model}"
+        if output_path.exists():
+            continue
+
+        print(f"Downloading {model} to {output_path}")
+        ultralytics.utils.downloads.attempt_download_asset(str(output_path))
+
+
+if __name__ == "__main__":
+    main()

install/setup.sh

@@ -1,15 +1,16 @@
 #!/bin/bash
 
-# Installs dependencies for building and running closed-set models
-# Usage: ./setup.sh [--no-models] [--no-closed-set] [--no-cuda]"
+# Installs dependencies for building and running semantic_inference
+# Usage: ./setup.sh [--no-models] [--no-closed-set] [--no-cuda] [--no-python]"
 #
-# This installs a minimal set of dependencies that *should* work, but YMMV.
+# This installs a minimal set of dependencies that *should* work for TensorRT, but YMMV.
 # Note: CUDA is hard-coded to install the 24.04 keyring and repos.
 #
 # Options:
 #   --no-models: Do not download model weights
 #   --no-closed-set: Do not install closed-set dependencies
 #   --no-cuda: Do not install the CUDA keyring (if you already have CUDA set up)
+#   --no-python: Do not set up python virtual environment for instance and open-set segmentation
 #   -h/--help: Show brief usage information
 
 function download_model() {
@@ -19,9 +20,14 @@ function download_model() {
     fi
 }
 
+
+# from https://stackoverflow.com/a/246128
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
 MODELS=true  # Download models
 CLOSED_SET=true  # Install closed-set deps
 CUDA=true  # Install cuda key
+MAKE_ENV=true  # Make python environment
 while [[ $# -gt 0 ]]; do
     case $1 in
         --no-models)
@@ -33,8 +39,11 @@ while [[ $# -gt 0 ]]; do
         --no-cuda)
             CUDA=false
             ;;
+        --no-python)
+            MAKE_ENV=false
+            ;;
         --help|-h)
-            echo "Usage: ./setup.sh [--no-models] [--no-closed-set] [--no-cuda]"
+            echo "Usage: ./setup.sh [--no-models] [--no-closed-set] [--no-cuda] [--no-python]"
             exit 0
     esac
     shift
@@ -61,9 +70,24 @@ if [ "$CLOSED_SET" = true ]; then
     sudo apt install libnvinfer-dev libnvonnxparsers-dev libnvinfer-plugin-dev "cuda-nvcc-${version}-${minor_version}" -y
 fi
 
+if [ "$MAKE_ENV" = true ]; then
+    PACKAGE_PATH="$(dirname ${SCRIPT_DIR})"
+    mkdir -p "$HOME"/.semantic_inference
+    python3 -m venv --system-site-packages --clear "$HOME"/.semantic_inference/env
+    source "$HOME"/.semantic_inference/env/bin/activate
+    python3 -m pip install -e "${PACKAGE_PATH}/semantic_inference"[openset]
+    deactivate
+fi
+
 if [ "$MODELS" = true ]; then
     mkdir -p "$HOME"/.semantic_inference/
     download_model ade20k-efficientvit_seg_l2.onnx "https://www.dropbox.com/scl/fi/qtaqm3htsdjlnoyqjvrol/ade20k-efficientvit_seg_l2.onnx?rlkey=3evg4gfeybd0wie0gom8535zo&st=1ocu1vl7&dl=1"
     download_model ade20k-hrnetv2-c1.onnx "https://www.dropbox.com/scl/fi/6hvyfgtk1j0uqt1t7dal9/ade20k-hrnetv2-c1.onnx?rlkey=k009j60g556h9p79bdoxgskbx&st=sqz7sqyl&dl=1"
     download_model ade20k-mobilenetv2dilated-c1_deepsup.onnx "https://www.dropbox.com/scl/fi/dlcizojblgaq8dnhd94o4/ade20k-mobilenetv2dilated-c1_deepsup.onnx?rlkey=5m7tv9x8bt0gsg1q77tert9ic&st=ixu66280&dl=1"
+    if [[ -e $HOME/.semantic_inference/env ]]; then
+        "$HOME"/.semantic_inference/env/bin/python3 "${SCRIPT_DIR}/download_weights.py" \
+            yolo11n-seg.pt \
+            yoloe-26m-seg.pt \
+            mobileclip2_b.ts
+    fi
 fi

semantic_inference/python/semantic_inference/__init__.py

@@ -31,7 +31,7 @@
 
 import pathlib
 
-from semantic_inference.misc import Logger
+from semantic_inference.misc import Logger, models_path
 
 
 def root_path():

semantic_inference/python/semantic_inference/misc.py

@@ -30,5 +30,15 @@
 """Various useful utilities."""
 
 import logging
+import os
+import pathlib
 
 Logger = logging.getLogger("semantic_inference")
+
+
+def models_path():
+    """Get path to ~/.semantic_inference directory."""
+    model_dir = os.getenv("SEMANTIC_INFERENCE_MODEL_DIR")
+    mpath = pathlib.Path(model_dir or "~/.semantic_inference").expanduser().absolute()
+    Logger.debug(f"Using model path: {mpath}")
+    return mpath

semantic_inference/python/semantic_inference/models/instance_segmenter.py

@@ -115,7 +115,7 @@ def segment(self, rgb_img, is_rgb_order=True):
             Encoded image
         """
         img = rgb_img if is_rgb_order else rgb_img[:, :, ::-1].copy()
-        return self._rotator.derotate(img)
+        return self(img)
 
     @property
     def device(self):
@@ -140,14 +140,14 @@ def forward(self, rgb_img):
         rotated = self._rotator.rotate(rgb_img)
         categories, masks, boxes, confidences = self.segmenter(rotated)
 
-        if self.masks is None:
+        if masks is None:
             instances = np.zeros(rgb_img.shape[:2])
         else:
             instances = np.zeros(masks[0].shape, dtype=np.uint32)
-            masks = self.masks.cpu().numpy()
-            category_ids = self.categories.cpu().numpy()
+            masks = masks.cpu().numpy()
+            category_ids = categories.cpu().numpy()
             for i in range(masks.shape[0]):
-                category_id = int(category_ids[i])  # category id are 0-indexed
+                category_id = int(category_ids[i]) + 1  # category id are 1-indexed
                 instance_id = i + 1  # instance ids are 1-indexed
                 # combine into single uint32
                 combined_id = (category_id << 16) | instance_id

semantic_inference/python/semantic_inference/models/wrappers.py

@@ -30,7 +30,6 @@
 """Model wrappers for image segmentation."""
 
 import dataclasses
-import os
 import pathlib
 
 import cv2
@@ -42,15 +41,7 @@
 from spark_config import Config, register_config
 from torchvision.ops import box_convert
 
-import semantic_inference.misc as misc
-
-
-def models_path():
-    """Get path to ~/.semantic_inference directory."""
-    model_dir = os.getenv("SEMANTIC_INFERENCE_MODEL_DIR")
-    mpath = pathlib.Path(model_dir or "~/.semantic_inference").expanduser().absolute()
-    misc.Logger.debug(f"Using model path: {mpath}")
-    return mpath
+from semantic_inference.misc import models_path
 
 
 class FastSAMSegmentation(nn.Module):
@@ -346,7 +337,8 @@ def load(cls, filepath):
 
 
 class YoloInstanceSegmenterBase(nn.Module):
-    """Base class for YOLO instance segmentation wrappers.
+    """
+    Base class for YOLO instance segmentation wrappers.
 
     Subclasses must assign self.model in their __init__.
     """
@@ -367,7 +359,7 @@ def eval(self):
     @property
     def category_names(self):
         """Get category names."""
-        return self.model.names
+        return [v for k, v in self.model.names.items()]
 
     def forward(self, img):
         """Segment image."""

semantic_inference/python/semantic_inference/visualization.py

@@ -327,8 +327,7 @@ def get_semantic_overlay_img(category_names, ret, img):
     masks = ret.masks
     boxes = ret.boxes
     confidences = ret.confidences
-
-    if masks is None:
+    if masks is None or len(masks.shape) == 2:
         # no bounding boxes to draw
         return img
 
@@ -341,15 +340,15 @@ def get_semantic_overlay_img(category_names, ret, img):
 
     # Overlay segmentation masks
     if masks is not None:
-        for i, mask_tensor in enumerate(masks.data):
+        for i, mask_tensor in enumerate(masks):
             box = boxes[i]
-            cls = int(categories[i].cpu().numpy())
+            cls = int(categories[i])
 
             # Get color for the class
             color = colors[cls].tolist()
 
             # Get mask and resize it to the image dimensions
-            mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
+            mask_np = mask_tensor.astype(np.uint8)
             mask_resized = cv2.resize(
                 mask_np,
                 (vis_img_bgr.shape[1], vis_img_bgr.shape[0]),
@@ -371,9 +370,9 @@ def get_semantic_overlay_img(category_names, ret, img):
 
     # Draw bounding boxes and labels
     for i, box in enumerate(boxes):
-        x1, y1, x2, y2 = map(int, box.cpu().numpy())
-        conf = confidences[i].cpu().numpy()
-        cls = int(categories[i].cpu().numpy())
+        x1, y1, x2, y2 = map(int, box)
+        conf = confidences[i]
+        cls = int(categories[i])
         label = f"{category_names[cls]} {conf:.2f}"
 
         color = colors[cls].tolist()

semantic_inference_msgs/CMakeLists.txt

@@ -20,6 +20,7 @@ rosidl_generate_interfaces(
   msg/FeatureVectors.msg
   msg/FeatureVectorStamped.msg
   msg/FeatureImage.msg
+  msg/Labelspace.msg
   srv/EncodeFeature.srv
   DEPENDENCIES
   builtin_interfaces

semantic_inference_msgs/msg/Labelspace.msg

@@ -0,0 +1,5 @@
+# Message containing labelspace information for a given model
+
+string[] names    # List of all label names (empty string if no name for specific label value)
+int32[] invalid   # All label IDs that should be ignored
+string metadata   # Additional labelspace metadata serialized via json

semantic_inference_ros/app/instance_segmentation_node

@@ -18,6 +18,7 @@ from std_msgs.msg import String
 import semantic_inference.models as models
 import semantic_inference_ros
 from semantic_inference.visualization import get_semantic_overlay_img
+from semantic_inference_msgs.msg import Labelspace
 from semantic_inference_ros import Conversions, ImageWorkerConfig
 
 
@@ -33,6 +34,12 @@ class InstanceSegmentationNodeConfig(sc.Config):
     nickname: str = "semantic_inference"
 
 
+def _latch_qos(depth):
+    return rclpy.qos.QoSProfile(
+        depth=depth, durability=rclpy.qos.DurabilityPolicy.TRANSIENT_LOCAL
+    )
+
+
 class InstanceSegmentationNode(Node):
     """Node to run instance segmentation."""
 
@@ -62,6 +69,10 @@ class InstanceSegmentationNode(Node):
         self._status_pub = self.create_publisher(String, "~/status", 1)
         self._hb_timer = self.create_timer(0.1, self._hb_callback)
 
+        qos = _latch_qos(1)
+        self._labelspace_pub = self.create_publisher(Labelspace, "labelspace", qos)
+        self._publish_labelspace()
+
         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
@@ -77,12 +88,11 @@ class InstanceSegmentationNode(Node):
         t0 = time.monotonic()
         with torch.no_grad():
             ret = self._model.segment(img, is_rgb_order=True)
-        self._inference_times_ms.append((time.monotonic() - t0) * 1e3)
 
+        self._inference_times_ms.append((time.monotonic() - t0) * 1e3)
         if ret.masks is None:
             self.get_logger().debug("No masks detected in the image.")
         else:
-            self.get_logger().debug("get image size: " + str(ret.img_shape))
             self.get_logger().debug(f"result category type: {ret.categories.dtype}")
 
         # Convert to int32 to match 32SC1 encoding expected by cv_bridge
@@ -99,6 +109,14 @@ class InstanceSegmentationNode(Node):
                 Conversions.to_image_msg(header, color_img, encoding="rgb8")
             )
 
+    def _publish_labelspace(self):
+        msg = Labelspace()
+        msg.names = [""] + self._model.category_names
+        # labels are 1-indexed and msg.names is 1 + len(model_categories)
+        metadata = {"object_labels": list(range(1, len(msg.names)))}
+        msg.metadata = json.dumps(metadata)
+        self._labelspace_pub.publish(msg)
+
     def _hb_callback(self):
         record = {
             "nickname": self._nickname,

semantic_inference_ros/launch/instance_segmentation.launch.yaml

@@ -1,19 +1,16 @@
 ---
 launch:
   # Segmentation configuration
-  - arg: {name: instance_seg_env, default: $(env HOME)/environments/gdsam2, description: Path to instance segmentation environment}
-  - arg: {name: model_name, default: gdsam2, description: Name of the segmentation model to use}
-  - arg: {name: config_path, default: $(find-pkg-share semantic_inference_ros)/config/instance_segmentation/$(var model_name).yaml, description: Configuration file for instance segmentation object detector}
-  # compression and logging
+  - arg: {name: env_path, default: $(env HOME)/.semantic_inference/env, description: Path to semantic_inference python environment}
+  - arg: {name: model_name, default: yoloe, description: Name of the segmentation model to use}
+  - arg: {name: config_path, default: $(find-pkg-share semantic_inference_ros)/config/instance_segmentation/$(var model_name).yaml, description: Configuration file for model}
+  # Node behavior
   - arg: {name: compressed_rgb, default: 'false', description: Triggers decompression for RGB stream}
-  # Run options
   - arg: {name: log-level, default: info, description: Set the ROS2 log level}
   - arg: {name: show_configs, default: 'true', description: Control logging config structs to stdout}
-  # Colormap and label grouping
   - arg: {name: run_backprojection, default: 'false', description: Triggers point cloud backprojection with RGB semantics}
   - arg: {name: backprojection_config, default: '{}', description: Extra parameters for backprojection}
-  # - arg: {name: colormap_path, default: $(find-pkg-share semantic_inference_ros)/config/distinct_150_colors.csv, description: Visualization colormap}
-  - arg: {name: colormap_path, default: $(find-pkg-share semantic_inference_ros)/config/hydra_object_gdsam_colors.csv, description: Visualization colormap}
+  - arg: {name: colormap_path, default: $(find-pkg-share semantic_inference_ros)/config/distinct_150_colors.csv, description: Visualization colormap}
   - arg: {name: image_topic, default: color/image_raw, description: Camera topic to use for backprojection coloring}
   # nodes
   - node:  # Decompress RGB stream
@@ -31,7 +28,7 @@ launch:
       pkg: semantic_inference_ros
       exec: instance_segmentation_node
       name: semantic_inference
-      pyenv: $(var instance_seg_env)
+      pyenv: $(var env_path)
       on_exit: shutdown
       param:
         - {name: config_path, value: $(var config_path), type: str}