-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #131 from WATonomous/lereljic/segformer-semantic-s…
…egmentation Segformer Semantic Segmentation
- Loading branch information
Showing
12 changed files
with
269 additions
and
25 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
30 changes: 30 additions & 0 deletions
30
src/perception/semantic_segmentation/launch/semantic_segmentation.launch.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
Empty file.
91 changes: 91 additions & 0 deletions
91
src/perception/semantic_segmentation/semantic_segmentation/segmentation_node.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,4 @@ | ||
| [develop] | ||
| script_dir=$base/lib/semantic_segmentation | ||
| [install] | ||
| install-scripts=$base/lib/semantic_segmentation |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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' | ||
| ], | ||
| }, | ||
| ) |