fix: allow for overriding pcl height / width.

common/helpers/cv2_pcl/src/cv2_pcl/__init__.py

@@ -3,18 +3,24 @@
 import ros_numpy as rnp
 import cv2
 
-from typing import Tuple
+from typing import Tuple, Union
 
 Mat = np.ndarray
 
 
-def pcl_to_cv2(pcl: PointCloud2) -> Mat:
+def pcl_to_cv2(
+    pcl: PointCloud2, height: Union[int, None] = None, width: Union[int, None] = None
+) -> Mat:
     """
     Convert a given PointCloud2 message to a cv2 image
     """
+
+    height = height or pcl.height
+    width = width or pcl.width
+
     # Extract rgb image from pointcloud
     frame = np.fromstring(pcl.data, dtype=np.uint8)
-    frame = frame.reshape(pcl.height, pcl.width, -1)
+    frame = frame.reshape(height, width, -1)
     frame = frame[:, :, 16:19]
 
     # Ensure array is contiguous
@@ -23,16 +29,25 @@ def pcl_to_cv2(pcl: PointCloud2) -> Mat:
     return frame
 
 
-def seg_to_centroid(pcl: PointCloud2, xyseg: np.ndarray) -> np.ndarray:
+def seg_to_centroid(
+    pcl: PointCloud2,
+    xyseg: np.ndarray,
+    height: Union[int, None] = None,
+    width: Union[int, None] = None,
+) -> np.ndarray:
     """
     Computes the centroid of a given segment in a pointcloud
     """
 
+    height = height or pcl.height
+    width = width or pcl.width
+
     # Convert xyseg to contours
     contours = xyseg.reshape(-1, 2)
 
     # cv2 image
-    cv_im = pcl_to_cv2(pcl)
+    cv_im = pcl_to_cv2(pcl, height, width)
+
     # Compute mask from contours
     mask = np.zeros(shape=cv_im.shape[:2])
     cv2.fillPoly(mask, pts=[contours], color=(255, 255, 255))
@@ -45,7 +60,7 @@ def seg_to_centroid(pcl: PointCloud2, xyseg: np.ndarray) -> np.ndarray:
 
     # Unpack pointcloud into xyz array
     pcl_xyz = rnp.point_cloud2.pointcloud2_to_xyz_array(pcl, remove_nans=False)
-    pcl_xyz = pcl_xyz.reshape(pcl.width, pcl.height, 3)
+    pcl_xyz = pcl_xyz.reshape(height, width, 3)
 
     # Extract points of interest
     xyz_points = [pcl_xyz[x][y] for x, y in indices]
@@ -54,17 +69,23 @@ def seg_to_centroid(pcl: PointCloud2, xyseg: np.ndarray) -> np.ndarray:
 
 
 def seg_to_minmax_z(
-    pcl: PointCloud2, xyseg: np.ndarray
+    pcl: PointCloud2,
+    xyseg: np.ndarray,
+    height: Union[int, None] = None,
+    width: Union[int, None] = None,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Computes the centroid of a given segment in a pointcloud
     """
 
+    height = height or pcl.height
+    width = width or pcl.width
+
     # Convert xyseg to contours
     contours = xyseg.reshape(-1, 2)
 
     # cv2 image
-    cv_im = pcl_to_cv2(pcl)
+    cv_im = pcl_to_cv2(pcl, height, width)
     # Compute mask from contours
     mask = np.zeros(shape=cv_im.shape[:2])
     cv2.fillPoly(mask, pts=[contours], color=(255, 255, 255))
@@ -77,7 +98,7 @@ def seg_to_minmax_z(
 
     # Unpack pointcloud into xyz array
     pcl_xyz = rnp.point_cloud2.pointcloud2_to_xyz_array(pcl, remove_nans=False)
-    pcl_xyz = pcl_xyz.reshape(pcl.height, pcl.width, 3)
+    pcl_xyz = pcl_xyz.reshape(height, width, 3)
 
     # Extract points of interest
     xyz_points = [pcl_xyz[x][y] for x, y in indices]
@@ -95,16 +116,27 @@ def seg_to_minmax_z(
     return centroid_min_z, centroid_max_z
 
 
-def bb_to_centroid(pcl: PointCloud2, x: int, y: int, w: int, h: int) -> np.ndarray:
+def bb_to_centroid(
+    pcl: PointCloud2,
+    x: int,
+    y: int,
+    w: int,
+    h: int,
+    height: Union[int, None] = None,
+    width: Union[int, None] = None,
+) -> np.ndarray:
     """
     Computes the centroid of a given bounding box in a pointcloud
     """
 
+    height = height or pcl.height
+    width = width or pcl.width
+
     # Convert xywh to bounding box coordinates.
     x1, y1, x2, y2 = x, y, x + w, y + h
 
     # cv2 image
-    frame = pcl_to_cv2(pcl)
+    frame = pcl_to_cv2(pcl, height, width)
     # Compute mask from bounding box
     mask = np.zeros(shape=frame.shape[:2])
     mask[y1:y2, x1:x2] = 255  # bounding box dimensions
@@ -116,6 +148,7 @@ def bb_to_centroid(pcl: PointCloud2, x: int, y: int, w: int, h: int) -> np.ndarr
 
     # Unpack pointcloud into xyz array
     pcl_xyz = rnp.point_cloud2.pointcloud2_to_xyz_array(pcl, remove_nans=False)
+    pcl_xyz = pcl_xyz.reshape(height, width, 3)
 
     # Extract points of interest
     xyz_points = [pcl_xyz[x][y] for x, y in indices]