From ec60b7645ba2f13422b17727cacbcdb5cf2d7c89 Mon Sep 17 00:00:00 2001
From: Florian Vahl <7vahl@informatik.uni-hamburg.de>
Date: Tue, 19 Mar 2024 12:51:08 +0100
Subject: [PATCH] Add distortion parameter

---
 ipm_image_node/ipm_image_node/ipm.py | 13 +++++++------
 ipm_library/ipm_library/ipm.py       |  9 +++++++--
 ipm_service/ipm_service/ipm.py       |  4 +++-
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/ipm_image_node/ipm_image_node/ipm.py b/ipm_image_node/ipm_image_node/ipm.py
index 5ad4569..85b55c8 100644
--- a/ipm_image_node/ipm_image_node/ipm.py
+++ b/ipm_image_node/ipm_image_node/ipm.py
@@ -35,17 +35,18 @@ class IPMImageNode(Node):
 
     def __init__(self) -> None:
         super().__init__('ipm_image_node')
+        # Declare params
+        self.declare_parameter('output_frame', 'base_footprint')
+        self.declare_parameter('type', 'mask')
+        self.declare_parameter('scale', 1.0)
+        self.declare_parameter('use_distortion', False)
+
         # We need to create a tf buffer
         self.tf_buffer = tf2.Buffer(cache_time=Duration(seconds=30.0))
         self.tf_listener = tf2.TransformListener(self.tf_buffer, self)
 
         # Create an IPM instance
-        self.ipm = IPM(self.tf_buffer)
-
-        # Declare params
-        self.declare_parameter('output_frame', 'base_footprint')
-        self.declare_parameter('type', 'mask')
-        self.declare_parameter('scale', 1.0)
+        self.ipm = IPM(self.tf_buffer, distortion=self.get_parameter('use_distortion').value)
 
         # Subscribe to camera info
         self.create_subscription(CameraInfo, 'camera_info', self.ipm.set_camera_info, 1)
diff --git a/ipm_library/ipm_library/ipm.py b/ipm_library/ipm_library/ipm.py
index e3dfb29..34c8586 100644
--- a/ipm_library/ipm_library/ipm.py
+++ b/ipm_library/ipm_library/ipm.py
@@ -35,7 +35,8 @@ class IPM:
     def __init__(
             self,
             tf_buffer: tf2_ros.Buffer,
-            camera_info: Optional[CameraInfo] = None) -> None:
+            camera_info: Optional[CameraInfo] = None,
+            distortion: bool = False) -> None:
         """
         Create a new inverse perspective mapper instance.
 
@@ -44,10 +45,13 @@ def __init__(
             camera intrinsics, camera frame, ...
             The camera info can be updated later on using the setter or
             provided directly if it is unlikly to change
+        :param distortion: Weather to use the distortion coefficients from the camera info.
+            Don't use this if you are using a rectified image.
         """
         # TF needs a listener that is init in the node context, so we need a reference
         self._tf_buffer = tf_buffer
         self.set_camera_info(camera_info)
+        self._distortion = distortion
 
     def set_camera_info(self, camera_info: CameraInfo) -> None:
         """
@@ -178,7 +182,8 @@ def map_points(
             self._camera_info,
             points,
             plane_normal,
-            plane_base_point)
+            plane_base_point,
+            use_distortion=self._distortion)
 
         # Transform output point if output frame if needed
         if output_frame_id not in [None, self._camera_info.header.frame_id]:
diff --git a/ipm_service/ipm_service/ipm.py b/ipm_service/ipm_service/ipm.py
index df0e792..542640a 100644
--- a/ipm_service/ipm_service/ipm.py
+++ b/ipm_service/ipm_service/ipm.py
@@ -28,11 +28,13 @@ class IPMService(Node):
 
     def __init__(self) -> None:
         super().__init__('ipm_service')
+        # Declare params
+        self.declare_parameter('use_distortion', False)
         # TF handling
         self.tf_buffer = tf2.Buffer(Duration(seconds=5))
         self.tf_listener = tf2.TransformListener(self.tf_buffer, self)
         # Create ipm library instance
-        self.ipm = IPM(self.tf_buffer)
+        self.ipm = IPM(self.tf_buffer, distortion=self.get_parameter('use_distortion').value)
         # Create subs
         self.camera_info_sub = self.create_subscription(
             CameraInfo, 'camera_info', self.ipm.set_camera_info, 1)