feat: add CUDA GPU acceleration to vision_node preprocessing
Move resize, BGR2GRAY, and GaussianBlur to GPU via cv2.cuda, with graceful CPU fallback when CUDA is unavailable.
This commit is contained in:
@@ -31,12 +31,45 @@ class VisionNode(Node):
|
|||||||
Point, '/vision/target', 10
|
Point, '/vision/target', 10
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# GPU acceleration setup
|
||||||
|
self._use_cuda = False
|
||||||
|
try:
|
||||||
|
if cv2.cuda.getCudaEnabledDeviceCount() > 0:
|
||||||
|
self._use_cuda = True
|
||||||
|
self._gpu_frame = cv2.cuda_GpuMat()
|
||||||
|
self._gpu_blurred = cv2.cuda_GpuMat()
|
||||||
|
self._gaussian_filter = cv2.cuda.createGaussianFilter(
|
||||||
|
cv2.CV_8UC1, cv2.CV_8UC1, (5, 5), 1)
|
||||||
|
except AttributeError:
|
||||||
|
pass
|
||||||
|
|
||||||
self.get_logger().info(
|
self.get_logger().info(
|
||||||
'VisionNode ready - subscribing to /camera/image_raw'
|
'VisionNode ready - %s, subscribing to /camera/image_raw'
|
||||||
|
% ('CUDA GPU acceleration enabled'
|
||||||
|
if self._use_cuda else 'CPU pipeline')
|
||||||
)
|
)
|
||||||
|
|
||||||
self.start_time = 0
|
self.start_time = 0
|
||||||
|
|
||||||
|
def _preprocess(self, cv_image):
|
||||||
|
"""Return blurred grayscale at half resolution."""
|
||||||
|
if self._use_cuda:
|
||||||
|
return self._gpu_preprocess(cv_image)
|
||||||
|
scaled = cv2.resize(cv_image, None, fx=0.5, fy=0.5,
|
||||||
|
interpolation=cv2.INTER_AREA)
|
||||||
|
gray = cv2.cvtColor(scaled, cv2.COLOR_BGR2GRAY)
|
||||||
|
return cv2.GaussianBlur(gray, (5, 5), 1)
|
||||||
|
|
||||||
|
def _gpu_preprocess(self, cv_image):
|
||||||
|
self._gpu_frame.upload(cv_image)
|
||||||
|
gpu_scaled = cv2.cuda.resize(
|
||||||
|
self._gpu_frame, None, fx=0.5, fy=0.5,
|
||||||
|
interpolation=cv2.INTER_AREA)
|
||||||
|
gpu_gray = cv2.cuda.cvtColor(gpu_scaled, cv2.COLOR_BGR2GRAY)
|
||||||
|
self._gpu_blurred.create(gpu_gray.size(), gpu_gray.type())
|
||||||
|
self._gaussian_filter.apply(gpu_gray, self._gpu_blurred)
|
||||||
|
return self._gpu_blurred.download()
|
||||||
|
|
||||||
def _image_callback(self, msg):
|
def _image_callback(self, msg):
|
||||||
try:
|
try:
|
||||||
encoding = msg.encoding
|
encoding = msg.encoding
|
||||||
@@ -46,21 +79,14 @@ class VisionNode(Node):
|
|||||||
|
|
||||||
cv_image = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
|
cv_image = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
|
||||||
|
|
||||||
# 分辨率缩小一半,提高处理速度
|
gauss_image = self._preprocess(cv_image)
|
||||||
scaled = cv2.resize(cv_image, None, fx=0.5, fy=0.5,
|
|
||||||
interpolation=cv2.INTER_AREA)
|
|
||||||
|
|
||||||
# 预处理:灰度图 + 高斯模糊 (去噪)
|
# Otsu 二值化
|
||||||
gray_image = cv2.cvtColor(scaled, cv2.COLOR_BGR2GRAY)
|
|
||||||
gauss_image = cv2.GaussianBlur(gray_image, (5, 5), 1)
|
|
||||||
|
|
||||||
# Otsu 二值化(替代 Canny,更稳健)
|
|
||||||
_, binary = cv2.threshold(
|
_, binary = cv2.threshold(
|
||||||
gauss_image, 0, 255,
|
gauss_image, 0, 255,
|
||||||
cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
|
cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||||
)
|
|
||||||
|
|
||||||
# 形态学闭运算:先膨胀后腐蚀,连接断裂区域、去除噪点
|
# 形态学闭运算
|
||||||
kernel = np.ones((3, 3), np.uint8)
|
kernel = np.ones((3, 3), np.uint8)
|
||||||
dilated = cv2.dilate(binary, kernel, iterations=1)
|
dilated = cv2.dilate(binary, kernel, iterations=1)
|
||||||
eroded = cv2.erode(dilated, kernel, iterations=1)
|
eroded = cv2.erode(dilated, kernel, iterations=1)
|
||||||
@@ -105,13 +131,13 @@ class VisionNode(Node):
|
|||||||
if moments['m00'] > 0:
|
if moments['m00'] > 0:
|
||||||
cx = moments['m10'] / moments['m00']
|
cx = moments['m10'] / moments['m00']
|
||||||
cy = moments['m01'] / moments['m00']
|
cy = moments['m01'] / moments['m00']
|
||||||
h, w = scaled.shape[:2]
|
h, w = gauss_image.shape[:2]
|
||||||
error_x = (cx - w / 2) / (w / 2)
|
error_x = (cx - w / 2) / (w / 2)
|
||||||
error_y = (cy - h / 2) / (h / 2)
|
error_y = (cy - h / 2) / (h / 2)
|
||||||
|
|
||||||
target_msg = Point()
|
target_msg = Point()
|
||||||
target_msg.x = error_x
|
target_msg.x = error_x
|
||||||
target_msg.y = error_y
|
target_msg.y = -error_y
|
||||||
target_msg.z = 0.0
|
target_msg.z = 0.0
|
||||||
self._target_publisher.publish(target_msg)
|
self._target_publisher.publish(target_msg)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user