Segformer Semantic Segmentation

docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile

@@ -1,13 +1,43 @@
-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=pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
+# ################################ Build library ################################
 
-################################ Source ################################
-FROM ${BASE_IMAGE} as source
+FROM ${BASE_PYTORCH_IMAGE} as Segformer
+
+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
+
+
+RUN apt-get update && apt-get install -y git ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6 libgl1-mesa-dev  \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 
 WORKDIR ${AMENT_WS}/src
+# Install MMSegmentation
+RUN git clone -b main https://github.com/open-mmlab/mmsegmentation.git /mmsegmentation
+WORKDIR /mmsegmentation
+ENV FORCE_CUDA="1"
+RUN pip install mmengine
+RUN pip install --no-cache-dir -e .
+RUN ["/bin/bash", "-c", "pip install openmim"]
+
+# Scan for dependencies
+RUN echo "openmim\nmmengine\nmmcv==2.0.1\n" > /tmp/pip_install_list.txt
+
+# Install segformer model
+RUN mim download mmsegmentation --config segformer_mit-b2_8xb1-160k_cityscapes-1024x1024 --dest ./model
+
+# # ################################ Source ################################
+
 
-# Copy in source code 
+FROM ${BASE_BUILD_IMAGE} as source
+WORKDIR ${AMENT_WS}/src
+
+# # 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 +46,36 @@ 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 \
+    && 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
+RUN pip3 install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
+RUN pip install cython
+RUN ["/bin/bash", "-c", "pip install 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 libopencv-dev
 
 # 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 +89,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 +101,37 @@ 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
+
+WORKDIR /mmsegmentation
+COPY --from=Segformer /mmsegmentation /mmsegmentation
+RUN pip install -r requirements.txt
+RUN pip install --no-cache-dir -e .
+WORKDIR ${AMENT_WS}
+# 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,4 @@
+semantic_segmentation_node:
+  ros__parameters:
+    input_topic: "/CAM_FRONT/image_rect_compressed"
+    publish_topic: "/camera/left/segmentations"

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

@@ -0,0 +1,24 @@
+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'
+    )
+    semantic_segmentation_node = Node(
+        package='semantic_segmentation',
+        executable='semantic_segmentation_node',
+        name='semantic_segmentation_node',
+        parameters=[config]
+    )
+
+    # 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,142 @@
+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 std_msgs.msg import String
+from cv_bridge import CvBridge
+from rich.logging import RichHandler
+import logging
+from ament_index_python.packages import get_package_share_directory
+# Configure the root logger
+logging.basicConfig(level=logging.INFO, handlers=[RichHandler(level=logging.WARNING)])
+
+
+# ALGO_VERSION = os.getenv("MODEL_NAME")
+
+# if not ALGO_VERSION:
+#     ALGO_VERSION = 'nvidia/segformer-b2-finetuned-ade-512-512'
+
+package_name = 'semantic_segmentation'
+package_share_directory = get_package_share_directory(package_name)
+CONFIG = os.path.join(package_share_directory, 'resource', 'model',
+                      'segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py')
+CHECKPOINT = os.path.join(package_share_directory, 'resource', 'model',
+                          'segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth')
+IMAGE_H = 900
+IMAGE_W = 1600
+
+# Adjust logging
+logging.getLogger('mmseg.apis').setLevel(logging.WARNING)
+logging.getLogger('mmengine').setLevel(logging.WARNING)
+
+"""Cityscapes palette that maps each class to RGB values."""
+COLOR_PALLETE = [
+    [128, 64, 128],  # road
+    [244, 35, 232],  # sidewalk
+    [70, 70, 70],    # building
+    [102, 102, 156],  # wall
+    [190, 153, 153],  # fence
+    [153, 153, 153],  # pole
+    [250, 170, 30],  # traffic light
+    [220, 220, 0],   # traffic sign
+    [107, 142, 35],  # vegetation
+    [152, 251, 152],  # terrain
+    [0, 130, 180],   # sky
+    [220, 20, 60],   # person
+    [255, 0, 0],     # rider
+    [0, 0, 142],     # car
+    [0, 0, 70],      # truck
+    [0, 60, 100],    # bus
+    [0, 80, 100],    # train
+    [0, 0, 230],     # motorcycle
+    [119, 11, 32],   # bicycle
+    [0, 0, 0]        # ignore/unlabeled
+]
+
+
+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(
+            'checkpoint', "model/segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth")
+
+        self.compressed = self.get_parameter('compressed').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(COLOR_PALLETE, dtype=np.uint8)
+        self.model = MMSegInferencer(CONFIG, 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, (1024, 1024))
+            else:
+                try:
+                    cv_image = self.cv_bridge.imgmsg_to_cv2(image, desired_encoding="passthrough")
+                    image = cv2.resize(cv_image, (1024, 1024))
+                except CvBridgeError as e:
+                    self.get_logger().error(str(e))
+                    return
+        with torch.no_grad():
+            out_img = self.model(image, show=False)['predictions']
+        # logits = torch.tensor(
+        # out_img, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        # Add batch and channel dimensions
+        # upsampled_logits = torch.nn.functional.interpolate(logits,
+        #                    size=(IMAGE_H, IMAGE_W),  # (height, width)
+        #                                                    mode='bilinear',
+        #                                                    align_corners=False)
+        # upsampled_logits = upsampled_logits.squeeze().numpy().astype(np.uint8)
+
+        #     color_seg[out_img == label, :] = color
+        color_seg = self.palette[out_img]
+        # img = np_image * 0.5 + color_seg * 0.5
+        # img_output = bridge.cv2_to_imgmsg(img)
+        # print(f'Publishing Segmentation')
+        # self.image_publisher.publish(img_output)
+        color_seg = cv2.resize(color_seg, (IMAGE_W, IMAGE_H))
+        mask_output = self.bridge.cv2_to_imgmsg(color_seg)
+        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