diff --git a/docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile b/docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile index 763cb632..e13516ae 100644 --- a/docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile +++ b/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//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/* diff --git a/modules/dev_overrides/docker-compose.perception.yaml b/modules/dev_overrides/docker-compose.perception.yaml index 4dcf2693..73fc96ac 100644 --- a/modules/dev_overrides/docker-compose.perception.yaml +++ b/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 diff --git a/modules/docker-compose.perception.yaml b/modules/docker-compose.perception.yaml index f095ff37..99dc9a49 100644 --- a/modules/docker-compose.perception.yaml +++ b/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: diff --git a/src/perception/semantic_segmentation/CMakeLists.txt b/src/perception/semantic_segmentation/CMakeLists.txt deleted file mode 100644 index 01396d95..00000000 --- a/src/perception/semantic_segmentation/CMakeLists.txt +++ /dev/null @@ -1,14 +0,0 @@ -cmake_minimum_required(VERSION 3.8) -project(semantic_segmentation) - -if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") - add_compile_options(-Wall -Wextra -Wpedantic) -endif() - -# find dependencies -find_package(ament_cmake REQUIRED) -# uncomment the following section in order to fill in -# further dependencies manually. -# find_package( REQUIRED) - -ament_package() diff --git a/src/perception/semantic_segmentation/config/params.yaml b/src/perception/semantic_segmentation/config/params.yaml new file mode 100644 index 00000000..236d15ee --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/perception/semantic_segmentation/launch/semantic_segmentation.launch.py b/src/perception/semantic_segmentation/launch/semantic_segmentation.launch.py new file mode 100644 index 00000000..938d4e1c --- /dev/null +++ b/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 diff --git a/src/perception/semantic_segmentation/package.xml b/src/perception/semantic_segmentation/package.xml index f90d0c8c..55a96e2a 100644 --- a/src/perception/semantic_segmentation/package.xml +++ b/src/perception/semantic_segmentation/package.xml @@ -3,14 +3,24 @@ semantic_segmentation 0.0.0 - TODO: Package description - bolty + The semantic segmentation package + Lucas Reljic TODO: License declaration - ament_cmake - + ament_copyright + ament_flake8 + ament_pep257 + python3-pytest + ament_lint_auto + ament_lint_common + sensor_msgs + geometry_msgs + std_msgs + OpenCV + cv_bridge + image_transport - ament_cmake + ament_python diff --git a/src/perception/semantic_segmentation/resource/semantic_segmentation b/src/perception/semantic_segmentation/resource/semantic_segmentation new file mode 100644 index 00000000..e69de29b diff --git a/src/perception/semantic_segmentation/semantic_segmentation/__init__.py b/src/perception/semantic_segmentation/semantic_segmentation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/perception/semantic_segmentation/semantic_segmentation/segmentation_node.py b/src/perception/semantic_segmentation/semantic_segmentation/segmentation_node.py new file mode 100644 index 00000000..7bbbc7cb --- /dev/null +++ b/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() diff --git a/src/perception/semantic_segmentation/setup.cfg b/src/perception/semantic_segmentation/setup.cfg new file mode 100644 index 00000000..790ce530 --- /dev/null +++ b/src/perception/semantic_segmentation/setup.cfg @@ -0,0 +1,4 @@ +[develop] +script_dir=$base/lib/semantic_segmentation +[install] +install-scripts=$base/lib/semantic_segmentation diff --git a/src/perception/semantic_segmentation/setup.py b/src/perception/semantic_segmentation/setup.py new file mode 100755 index 00000000..4eb5f41a --- /dev/null +++ b/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='lereljic@watonomous.ca', + description='TODO: Package description', + license='TODO: License declaration', + tests_require=['pytest'], + entry_points={ + 'console_scripts': [ + 'semantic_segmentation_node = semantic_segmentation.segmentation_node:main' + ], + }, +)