Segformer Semantic Segmentation

docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile

@@ -1,13 +1,20 @@
-ARG BASE_IMAGE=ghcr.io/watonomous/wato_monorepo/base:humble-ubuntu22.04
+ARG BASE_BUILD_IMAGE=ghcr.io/watonomous/wato_monorepo/base:cuda11.7-humble-ubuntu22.04-devel
+ARG BASE_PROD_IMAGE=ghcr.io/watonomous/wato_monorepo/base:cuda11.7-humble-ubuntu22.04
+ARG BASE_PYTORCH_IMAGE=ghcr.io/watonomous/wato_monorepo/segformer_segmentation:latest
+# ################################ Build library ################################
 
-################################ Source ################################
-FROM ${BASE_IMAGE} as source
+FROM ${BASE_PYTORCH_IMAGE} as Segformer
 
+# # ################################ Source ################################
+
+
+FROM ${BASE_BUILD_IMAGE} as source
 WORKDIR ${AMENT_WS}/src
 
-# Copy in source code 
+# # Copy in source code 
 COPY src/perception/semantic_segmentation semantic_segmentation
 COPY src/wato_msgs/sample_msgs sample_msgs
+COPY --from=Segformer /mmsegmentation/model ${AMENT_WS}/src/semantic_segmentation/resource/model
 
 # Scan for rosdeps
 RUN apt-get -qq update && rosdep update && \
@@ -16,12 +23,42 @@ RUN apt-get -qq update && rosdep update && \
         | awk '{print $3}' \
         | sort  > /tmp/colcon_install_list
 
+
 ################################# Dependencies ################################
-FROM ${BASE_IMAGE} as dependencies
+FROM ${BASE_BUILD_IMAGE} as dependencies
+
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    python3-pip \
+    ninja-build \
+    libglib2.0-0 \
+    libsm6 \
+    libxrender-dev \
+    libxext6 \
+    libgl1-mesa-dev \
+    libopencv-dev \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Segformer dependencies 
+# COPY --from=Segformer /tmp/pip_install_list.txt /tmp/pip_install_list.txt torchaudio==0.13.1
+RUN pip install torch==1.13.1+cu116 \
+    torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 \
+    cython \
+    https://download.openmmlab.com/mmcv/dist/cu117/torch1.13.0/mmcv-2.0.0rc4-cp310-cp310-manylinux1_x86_64.whl
+RUN apt update && apt install -y ros-humble-cv-bridge
+
+WORKDIR /mmsegmentation
+COPY --from=Segformer /mmsegmentation /mmsegmentation
+RUN pip install -r requirements.txt --no-cache-dir -e . && pip uninstall numpy -y && pip install numpy==1.26.4
 
 # Install Rosdep requirements
 COPY --from=source /tmp/colcon_install_list /tmp/colcon_install_list
-RUN apt-fast install -qq -y --no-install-recommends $(cat /tmp/colcon_install_list)
+RUN apt install -qq -y --no-install-recommends $(cat /tmp/colcon_install_list)
 
 # Copy in source code from source stage
 WORKDIR ${AMENT_WS}
@@ -35,6 +72,10 @@ RUN apt-get -qq autoremove -y && apt-get -qq autoclean && apt-get -qq clean && \
 ################################ Build ################################
 FROM dependencies as build
 
+
+ENV FORCE_CUDA="1"
+
+
 # Build ROS2 packages
 WORKDIR ${AMENT_WS}
 RUN . /opt/ros/$ROS_DISTRO/setup.sh && \
@@ -43,11 +84,32 @@ RUN . /opt/ros/$ROS_DISTRO/setup.sh && \
 
 # Entrypoint will run before any CMD on launch. Sources ~/opt/<ROS_DISTRO>/setup.bash and ~/ament_ws/install/setup.bash
 COPY docker/wato_ros_entrypoint.sh ${AMENT_WS}/wato_ros_entrypoint.sh
+
+# Add runtime libraries to path
+ENV CUDNN_DIR=/mmsegmentation/cuda
+ENV CV2_CUDABACKEND=0
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+
 ENTRYPOINT ["./wato_ros_entrypoint.sh"]
 
 ################################ Prod ################################
 FROM build as deploy
 
+# Install runtime libs
+RUN apt-get update && apt-get install -y \
+    ros-humble-cv-bridge
+
+
+WORKDIR ${AMENT_WS}
+
+RUN mkdir -p install/semantic_segmentation/lib/
+# Add runtime libraries to path
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+
+COPY docker/wato_ros_entrypoint.sh ${AMENT_WS}/wato_ros_entrypoint.sh
+ENTRYPOINT ["./wato_ros_entrypoint.sh"]
 # Source Cleanup and Security Setup
 RUN chown -R $USER:$USER ${AMENT_WS}
 RUN rm -rf src/*

modules/dev_overrides/docker-compose.perception.yaml

@@ -44,6 +44,7 @@ services:
     command: tail -F anything
     volumes:
       - ${MONO_DIR}/src/perception/semantic_segmentation:/home/bolty/ament_ws/src/semantic_segmentation
+      - /mnt/wato-drive/perception/segformer-b2:/home/bolty/ament_ws/src/semantic_segmentation/resource/model
 
   lane_detection:
     <<: *fixuid

modules/docker-compose.perception.yaml

@@ -56,6 +56,16 @@ services:
       target: deploy
     image: "${PERCEPTION_SEMANTIC_SEGMENTATION_IMAGE}:${TAG}"
     command: /bin/bash -c "ros2 launch semantic_segmentation semantic_segmentation.launch.py"
+    volumes: 
+      - /mnt/wato-drive/perception/segformer-b2:/home/bolty/ament_ws/src/semantic_segmentation/resource/model
+          # add gpus all
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [gpu]
+              count: 1
 
   lane_detection:
     build:

src/perception/semantic_segmentation/config/params.yaml

@@ -0,0 +1,8 @@
+semantic_segmentation_node:
+  ros__parameters:
+    input_topic: "/CAM_FRONT/image_rect_compressed"
+    publish_topic: "/camera/left/segmentations"
+    config: "model/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py"
+    checkpoint: "model/segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth"
+    MODEL_IMAGE_H: 1024
+    MODEL_IMAGE_W: 1024

src/perception/semantic_segmentation/launch/semantic_segmentation.launch.py

@@ -0,0 +1,30 @@
+from launch import LaunchDescription
+from launch_ros.actions import Node
+from ament_index_python.packages import get_package_share_directory
+import os
+
+
+def generate_launch_description():
+    ld = LaunchDescription()
+    config = os.path.join(
+        get_package_share_directory('semantic_segmentation'),
+        'config',
+        'params.yaml'
+    )
+
+    resource_path = os.path.join(
+        get_package_share_directory('semantic_segmentation'),
+        'resource'
+    )
+
+    semantic_segmentation_node = Node(
+        package='semantic_segmentation',
+        executable='semantic_segmentation_node',
+        name='semantic_segmentation_node',
+        parameters=[config,
+                    {'resource_path': resource_path}]
+    )
+
+    # finalize
+    ld.add_action(semantic_segmentation_node)
+    return ld

src/perception/semantic_segmentation/package.xml

@@ -3,14 +3,24 @@
 <package format="3">
   <name>semantic_segmentation</name>
   <version>0.0.0</version>
-  <description>TODO: Package description</description>
-  <maintainer email="e23zhou@uwaterloo.ca">bolty</maintainer>
+  <description>The semantic segmentation package</description>
+  <maintainer email="lereljic@watonomous.ca">Lucas Reljic</maintainer>
   <license>TODO: License declaration</license>
 
-  <buildtool_depend>ament_cmake</buildtool_depend>
-
+  <test_depend>ament_copyright</test_depend>
+  <test_depend>ament_flake8</test_depend>
+  <test_depend>ament_pep257</test_depend>
+  <test_depend>python3-pytest</test_depend>
+  <test_depend>ament_lint_auto</test_depend>
+  <test_depend>ament_lint_common</test_depend>
+  <depend>sensor_msgs</depend>
+  <depend>geometry_msgs</depend>
+  <depend>std_msgs</depend>
+  <depend>OpenCV</depend>
+  <depend>cv_bridge</depend>
+  <depend>image_transport</depend>
 
   <export>
-    <build_type>ament_cmake</build_type>
+    <build_type>ament_python</build_type>
   </export>
 </package>

src/perception/semantic_segmentation/semantic_segmentation/segmentation_node.py

@@ -0,0 +1,91 @@
+import numpy as np
+import cv2
+import os
+from mmseg.apis import MMSegInferencer
+import torch
+import rclpy
+from rclpy.node import Node
+from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSHistoryPolicy
+from sensor_msgs.msg import Image, CompressedImage
+from cv_bridge import CvBridge
+import logging
+
+
+class SemanticSegmentation(Node):
+    def __init__(self):
+        super().__init__('semantic_segmentation_node')
+        self.declare_parameter('pub_image', True)
+        self.declare_parameter('pub_masks', True)
+        self.declare_parameter('compressed', True)
+        self.declare_parameter('config', "model/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py")
+        self.declare_parameter('resource_path', "")
+        self.declare_parameter(
+            'checkpoint', "model/segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth")
+        self.declare_parameter('MODEL_IMAGE_H', 1024)
+        self.declare_parameter('MODEL_IMAGE_W', 1024)
+
+        self.config = os.path.join(self.get_parameter(
+            'resource_path').value, self.get_parameter('config').value)
+        self.checkpoint = os.path.join(self.get_parameter(
+            'resource_path').value, self.get_parameter('checkpoint').value)
+        self.compressed = self.get_parameter('compressed').value
+        self.modelH = self.get_parameter('MODEL_IMAGE_H').value
+        self.modelW = self.get_parameter('MODEL_IMAGE_W').value
+
+        self.image_subscription = self.create_subscription(
+            Image if not self.compressed else CompressedImage,
+            "/CAM_FRONT/image_rect_compressed",
+            self.listener_callback,
+            qos_profile=QoSProfile(
+                reliability=QoSReliabilityPolicy.RELIABLE,
+                history=QoSHistoryPolicy.KEEP_LAST,
+                depth=0,
+            ),
+        )
+
+        self.image_publisher = self.create_publisher(
+            Image,
+            '/camera/left/segmentations',
+            10
+        )
+        # self.palette = np.array(self.palette, dtype=np.uint8)
+        self.model = MMSegInferencer(self.config, self.checkpoint,
+                                     dataset_name="cityscapes", device='cuda:0')
+        self.bridge = CvBridge()
+
+    def listener_callback(self, msg):
+        images = [msg]  # msg is a single sensor image
+        for image in images:
+            # convert ros Image to cv::Mat
+            if self.compressed:
+                np_arr = np.frombuffer(msg.data, np.uint8)
+                cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+                image = cv2.resize(cv_image, (self.modelW, self.modelH))
+            else:
+                try:
+                    cv_image = self.cv_bridge.imgmsg_to_cv2(image, desired_encoding="passthrough")
+                    image = cv2.resize(cv_image, (self.modelW, self.modelH))
+                except CvBridgeError as e:
+                    self.get_logger().error(str(e))
+                    return
+        with torch.no_grad():
+            out_img = self.model(image, show=False)['predictions']
+
+        out_array = np.array(out_img, np.uint8)
+        mask_output = self.bridge.cv2_to_imgmsg(out_array)
+        self.image_publisher.publish(mask_output)
+
+
+def main(args=None):
+
+    rclpy.init(args=args)
+    semantic_segmentation_node = SemanticSegmentation()
+
+    rclpy.spin(semantic_segmentation_node)
+
+    semantic_segmentation_node.destroy_node()
+    rclpy.shutdown()
+
+
+if __name__ == '__main__':
+    main()

src/perception/semantic_segmentation/setup.cfg

@@ -0,0 +1,4 @@
+[develop]
+script_dir=$base/lib/semantic_segmentation
+[install]
+install-scripts=$base/lib/semantic_segmentation

src/perception/semantic_segmentation/setup.py

@@ -0,0 +1,42 @@
+import os
+from glob import glob
+from setuptools import setup
+
+package_name = 'semantic_segmentation'
+
+setup(
+    name=package_name,
+    version='0.0.0',
+    packages=[package_name],
+    data_files=[
+        ('share/ament_index/resource_index/packages',
+            ['resource/' + package_name]),
+        ('share/' + package_name, ['package.xml']),
+        ('share/' + package_name + '/resource/model', glob('resource/model/*')),
+        (os.path.join('share', package_name, 'launch'), glob('launch/*.py')),
+        (os.path.join('share', package_name, 'config'), glob('config/*.yaml')),
+    ],
+    install_requires=[
+        'setuptools',
+        'transformers',
+        'torch',
+        'cython',
+        'datasets',
+        "Pillow",
+        'charset-normalizer==2.0.0',
+        'packaging==20.9',
+        'numpy==1.23.0',
+        'opencv-python'
+    ],
+    zip_safe=True,
+    maintainer='Lucas',
+    maintainer_email='[email protected]',
+    description='TODO: Package description',
+    license='TODO: License declaration',
+    tests_require=['pytest'],
+    entry_points={
+        'console_scripts': [
+            'semantic_segmentation_node = semantic_segmentation.segmentation_node:main'
+        ],
+    },
+)

src/samples/python/aggregator/setup.py

@@ -14,7 +14,7 @@
         # Include our package.xml file
         (os.path.join('share', package_name), ['package.xml']),
         # Include all launch files.
-        (os.path.join('share', package_name, 'launch'), \
+        (os.path.join('share', package_name, 'launch'),
          glob(os.path.join('launch', '*.launch.py'))),
     ],
     install_requires=['setuptools'],