Real time semantic segmentation node (#77)

* Clean up dockerfile and docker compose

* Add space optimized prod image stage

* Clean up code

* Update image with correct header and create save image directories

* Add visualization topics

* Address PR comments, add perception_utils

---------

Co-authored-by: Steven Gong <[email protected]>
Co-authored-by: Eddy Zhou <[email protected]>

Author: 3 people

docker/perception/semantic_segmentation/semantic_segmentation.Dockerfile

@@ -1,20 +1,23 @@
-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 ################################
+ARG BASE_IMAGE=ghcr.io/watonomous/wato_monorepo/base:cuda12.0-humble-ubuntu22.04-devel
+ARG RUNTIME_IMAGE=ghcr.io/watonomous/wato_monorepo/base:cuda12.0-humble-ubuntu22.04
+ARG PADDLE_INFERENCE_BUILD_URL=ghcr.io/watonomous/perception_paddlepaddle_inference_build_cuda-12.0
+################################ Build library ################################
+FROM ${PADDLE_INFERENCE_BUILD_URL} as PADDLE_INFERENCE_BUILD
 
-FROM ${BASE_PYTORCH_IMAGE} as Segformer
+################################ Source ################################
+FROM ${BASE_IMAGE} as source
 
-# # ################################ Source ################################
-
-
-FROM ${BASE_BUILD_IMAGE} as source
 WORKDIR ${AMENT_WS}/src
 
-# # Copy in source code 
+# Copy in the paddle inference library
+RUN mkdir -p semantic_segmentation/src
+COPY --from=PADDLE_INFERENCE_BUILD /paddle/paddle_inference_cuda120_build.tar /paddle/paddle_inference_cuda120_build.tar
+RUN tar -xvf /paddle/paddle_inference_cuda120_build.tar -C /paddle/
+RUN rm /paddle/paddle_inference_cuda120_build.tar
+
+# 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
+COPY src/perception/perception_utils perception_utils
 
 # Scan for rosdeps
 RUN apt-get -qq update && rosdep update && \
@@ -25,43 +28,17 @@ RUN apt-get -qq update && rosdep update && \
 
 
 ################################# Dependencies ################################
-FROM ${BASE_BUILD_IMAGE} as dependencies
+FROM ${BASE_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
+RUN apt update && apt install -y tensorrt ros-humble-cv-bridge libopencv-dev
 
 # Install Rosdep requirements
 COPY --from=source /tmp/colcon_install_list /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}
+COPY --from=source /paddle /paddle
 COPY --from=source ${AMENT_WS}/src src
 
 # Dependency Cleanup
@@ -72,10 +49,6 @@ 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 && \
@@ -86,30 +59,43 @@ RUN . /opt/ros/$ROS_DISTRO/setup.sh && \
 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
+# ################################ Prod ################################
+FROM ${RUNTIME_IMAGE} as deploy
+
 
 # Install runtime libs
 RUN apt-get update && apt-get install -y \
-    ros-humble-cv-bridge
+    ros-humble-cv-bridge \
+    tensorrt
 
+# Copy the compiled binary to the runtime image
+COPY --from=build ${AMENT_WS} ${AMENT_WS}
 
 WORKDIR ${AMENT_WS}
 
+# Copy in the paddle inference library
 RUN mkdir -p install/semantic_segmentation/lib/
+COPY --from=PADDLE_INFERENCE_BUILD /paddle/paddle_inference_cuda120_build.tar install/semantic_segmentation/lib/paddle_inference_cuda120_build.tar
+RUN tar -xvf install/semantic_segmentation/lib/paddle_inference_cuda120_build.tar -C install/semantic_segmentation/lib
+RUN rm install/semantic_segmentation/lib/paddle_inference_cuda120_build.tar
+
 # 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
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/paddle_inference_cuda120_build/paddle_inference_install_dir/paddle/lib/
+
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/paddle_inference_cuda120_build/paddle_inference_install_dir/third_party/install/cryptopp/lib/
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/paddle_inference_cuda120_build/paddle_inference_install_dir/third_party/install/mkldnn/lib/
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${AMENT_WS}/install/semantic_segmentation/lib/paddle_inference_cuda120_build/paddle_inference_install_dir/third_party/install/mklml/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/docker-compose.perception.yaml

@@ -47,6 +47,7 @@ services:
     command: /bin/bash -c "ros2 launch lidar_object_detection nuscenes_launch.py"
     volumes:
       - /mnt/wato-drive2/perception_models/transfusion_trained_model.pth:/home/bolty/OpenPCDet/models/transfusion_trained_model.pth
+
   semantic_segmentation:
     build:
       context: ..
@@ -56,17 +57,16 @@ services:
         - "${PERCEPTION_SEMANTIC_SEGMENTATION_IMAGE}:build_main"
       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
+    command: /bin/bash -c "ros2 launch semantic_segmentation eve.launch.py"
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              capabilities: [gpu]
               count: 1
+              capabilities: [ gpu ]
+    volumes: 
+      - /mnt/wato-drive2/perception-weights/semantic_segmentation/pp_liteseg_infer_model:/perception_models/semantic_segmentation/pp_liteseg_infer_model
 
   lane_detection:
     build:
@@ -108,4 +108,4 @@ services:
         - "${PERCEPTION_DEPTH_ESTIMATION_IMAGE}:build_main"
       target: deploy
     image: "${PERCEPTION_DEPTH_ESTIMATION_IMAGE}:${TAG}"
-    command: /bin/bash -c "ros2 launch depth_estimation eve.launch.py"
+    command: /bin/bash -c "ros2 launch depth_estimation eve.launch.py"

src/perception/perception_utils/CMakeLists.txt

@@ -0,0 +1,49 @@
+cmake_minimum_required(VERSION 3.5)
+project(perception_utils)
+
+# Default to C++14
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 14)
+endif()
+
+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)
+find_package(rclcpp REQUIRED)
+find_package(OpenCV REQUIRED)
+
+# Include header files
+include_directories(include
+  ${OpenCV_INCLUDE_DIRS}
+)
+
+# Declare a C++ library
+add_library(perception_utils SHARED
+  src/camera_utils.cpp
+)
+
+# Specify libraries to link a library or executable target against
+ament_target_dependencies(perception_utils
+  rclcpp
+  OpenCV)
+
+install(TARGETS perception_utils
+  EXPORT export_perception_utils
+  ARCHIVE DESTINATION lib
+  LIBRARY DESTINATION lib
+  RUNTIME DESTINATION bin
+)
+
+ament_export_targets(export_perception_utils HAS_LIBRARY_TARGET)
+ament_export_include_directories(include)
+ament_export_dependencies(rclcpp)
+
+# Install header files
+install(DIRECTORY include/
+  DESTINATION include
+)
+
+ament_package()

src/perception/perception_utils/include/perception_utils/camera_utils.hpp

@@ -0,0 +1,14 @@
+#ifndef CAMERA_UTILS_HPP_
+#define CAMERA_UTILS_HPP_
+
+#include <opencv2/opencv.hpp>
+
+namespace CameraUtils {
+// Swaps the channels of an image from HWC to CHW format
+void hwc_img_2_chw_data(const cv::Mat& hwc_img, float* data);
+cv::Mat resize_image_aspect_ratio(const cv::Mat& original_image, int max_width, int max_height);
+cv::Mat resize_with_padding(const cv::Mat& original_image, int target_width, int target_height);
+cv::Mat resize_from_center(const cv::Mat& original_image, int target_width, int target_height);
+};  // namespace CameraUtils
+
+#endif  // CAMERA_UTILS_HPP_

src/perception/perception_utils/package.xml

@@ -0,0 +1,18 @@
+<?xml version="1.0"?>
+<package format="3">
+  <name>perception_utils</name>
+  <version>0.0.1</version>
+  <description>A utility library for perception tasks in ROS 2.</description>
+
+  <maintainer email="[email protected]">Justin Leung</maintainer>
+  <license>Apache-2.0</license>
+
+  <buildtool_depend>ament_cmake</buildtool_depend>
+  <depend>rclcpp</depend>
+  <depend>opencv</depend>
+
+  <export>
+    <build_type>ament_cmake</build_type>
+  </export>
+
+</package>

src/perception/perception_utils/src/camera_utils.cpp

@@ -0,0 +1,97 @@
+#include "perception_utils/camera_utils.hpp"
+#include <opencv2/opencv.hpp>
+
+// Swaps the channels of an image from HWC to CHW format
+void CameraUtils::hwc_img_2_chw_data(const cv::Mat& hwc_img, float* data) {
+  int rows = hwc_img.rows;
+  int cols = hwc_img.cols;
+  int chs = hwc_img.channels();
+  for (int i = 0; i < chs; ++i) {
+    cv::extractChannel(hwc_img, cv::Mat(rows, cols, CV_32FC1, data + i * rows * cols), i);
+  }
+}
+
+cv::Mat CameraUtils::resize_image_aspect_ratio(const cv::Mat& original_image, int max_width,
+                                               int max_height) {
+  int original_width = original_image.cols;
+  int original_height = original_image.rows;
+  double original_aspect_ratio = (double)original_width / original_height;
+
+  int new_width, new_height;
+  double max_aspect_ratio = (double)max_width / max_height;
+
+  if (original_aspect_ratio > max_aspect_ratio) {
+    // Width is the limiting factor
+    new_width = max_width;
+    new_height = static_cast<int>(max_width / original_aspect_ratio);
+  } else {
+    // Height is the limiting factor
+    new_height = max_height;
+    new_width = static_cast<int>(max_height * original_aspect_ratio);
+  }
+
+  cv::Mat resized_image;
+  cv::resize(original_image, resized_image, cv::Size(new_width, new_height));
+
+  return resized_image;
+}
+
+cv::Mat CameraUtils::resize_with_padding(const cv::Mat& original_image, int target_width,
+                                         int target_height) {
+  int original_width = original_image.cols;
+  int original_height = original_image.rows;
+
+  double target_ratio = (double)target_width / target_height;
+  double original_ratio = (double)original_width / original_height;
+
+  int new_width, new_height;
+
+  if (original_ratio > target_ratio) {
+    // Original is wider. Fit to width and pad height.
+    new_width = target_width;
+    new_height =
+        static_cast<int>(original_height * (static_cast<double>(target_width) / original_width));
+  } else {
+    // Original is taller. Fit to height and pad width.
+    new_height = target_height;
+    new_width =
+        static_cast<int>(original_width * (static_cast<double>(target_height) / original_height));
+  }
+
+  cv::Mat resized_image;
+  cv::resize(original_image, resized_image, cv::Size(new_width, new_height));
+
+  int top = (target_height - new_height) / 2;
+  int bottom = target_height - new_height - top;
+  int left = (target_width - new_width) / 2;
+  int right = target_width - new_width - left;
+
+  cv::Mat padded_image;
+  cv::copyMakeBorder(resized_image, padded_image, top, bottom, left, right, cv::BORDER_CONSTANT,
+                     cv::Scalar(0, 0, 0));
+
+  return padded_image;
+}
+
+cv::Mat CameraUtils::resize_from_center(const cv::Mat& original_image, int target_width,
+                                        int target_height) {
+  int original_width = original_image.cols;
+  int original_height = original_image.rows;
+
+  // Calculate the new height maintaining the aspect ratio
+  double target_ratio = (double)target_width / target_height;
+  int new_height = static_cast<int>(original_width / target_ratio);
+
+  // Calculate the cropping area
+  int startY = (original_height - new_height) / 2;
+
+  // Crop the image from the center
+  cv::Rect roi(0, startY, original_width, new_height);
+  cv::Mat cropped_image = original_image(roi);
+
+  // Resize the cropped image
+  cv::Mat resized_image;
+  cv::resize(cropped_image, resized_image, cv::Size(target_width, target_height));
+
+  return resized_image;
+}