diff --git a/common/helpers/cv2_pcl/src/cv2_pcl/__init__.py b/common/helpers/cv2_pcl/src/cv2_pcl/__init__.py index 470f29bad..3cbed3751 100644 --- a/common/helpers/cv2_pcl/src/cv2_pcl/__init__.py +++ b/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]