diff --git a/install/download_weights.py b/install/download_weights.py new file mode 100755 index 0000000..3cadf61 --- /dev/null +++ b/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() diff --git a/install/setup.sh b/install/setup.sh index 4a957cb..054a4c1 100755 --- a/install/setup.sh +++ b/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 diff --git a/semantic_inference/python/semantic_inference/__init__.py b/semantic_inference/python/semantic_inference/__init__.py index 8d53d73..3b9ba7b 100644 --- a/semantic_inference/python/semantic_inference/__init__.py +++ b/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(): diff --git a/semantic_inference/python/semantic_inference/misc.py b/semantic_inference/python/semantic_inference/misc.py index 1e9fc98..514e7de 100644 --- a/semantic_inference/python/semantic_inference/misc.py +++ b/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 diff --git a/semantic_inference/python/semantic_inference/models/instance_segmenter.py b/semantic_inference/python/semantic_inference/models/instance_segmenter.py index 44354b0..01937a9 100644 --- a/semantic_inference/python/semantic_inference/models/instance_segmenter.py +++ b/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 diff --git a/semantic_inference/python/semantic_inference/models/wrappers.py b/semantic_inference/python/semantic_inference/models/wrappers.py index a528b56..4ffc0df 100644 --- a/semantic_inference/python/semantic_inference/models/wrappers.py +++ b/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.""" diff --git a/semantic_inference/python/semantic_inference/visualization.py b/semantic_inference/python/semantic_inference/visualization.py index 44cc7ed..f730feb 100644 --- a/semantic_inference/python/semantic_inference/visualization.py +++ b/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() diff --git a/semantic_inference_msgs/CMakeLists.txt b/semantic_inference_msgs/CMakeLists.txt index 1618977..f9986ad 100644 --- a/semantic_inference_msgs/CMakeLists.txt +++ b/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 diff --git a/semantic_inference_msgs/msg/Labelspace.msg b/semantic_inference_msgs/msg/Labelspace.msg new file mode 100644 index 0000000..4614d7d --- /dev/null +++ b/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 diff --git a/semantic_inference_ros/app/instance_segmentation_node b/semantic_inference_ros/app/instance_segmentation_node index 9e6776e..30b44a1 100755 --- a/semantic_inference_ros/app/instance_segmentation_node +++ b/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, diff --git a/semantic_inference_ros/launch/instance_segmentation.launch.yaml b/semantic_inference_ros/launch/instance_segmentation.launch.yaml index 1856fe8..3e38cb7 100644 --- a/semantic_inference_ros/launch/instance_segmentation.launch.yaml +++ b/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}