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
|
||||
)
|
||||
|
||||
# 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(
|
||||
'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
|
||||
|
||||
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):
|
||||
try:
|
||||
encoding = msg.encoding
|
||||
@@ -46,21 +79,14 @@ class VisionNode(Node):
|
||||
|
||||
cv_image = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
|
||||
|
||||
# 分辨率缩小一半,提高处理速度
|
||||
scaled = cv2.resize(cv_image, None, fx=0.5, fy=0.5,
|
||||
interpolation=cv2.INTER_AREA)
|
||||
gauss_image = self._preprocess(cv_image)
|
||||
|
||||
# 预处理:灰度图 + 高斯模糊 (去噪)
|
||||
gray_image = cv2.cvtColor(scaled, cv2.COLOR_BGR2GRAY)
|
||||
gauss_image = cv2.GaussianBlur(gray_image, (5, 5), 1)
|
||||
|
||||
# Otsu 二值化(替代 Canny,更稳健)
|
||||
# Otsu 二值化
|
||||
_, binary = cv2.threshold(
|
||||
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)
|
||||
dilated = cv2.dilate(binary, kernel, iterations=1)
|
||||
eroded = cv2.erode(dilated, kernel, iterations=1)
|
||||
@@ -105,13 +131,13 @@ class VisionNode(Node):
|
||||
if moments['m00'] > 0:
|
||||
cx = moments['m10'] / 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_y = (cy - h / 2) / (h / 2)
|
||||
|
||||
target_msg = Point()
|
||||
target_msg.x = error_x
|
||||
target_msg.y = error_y
|
||||
target_msg.y = -error_y
|
||||
target_msg.z = 0.0
|
||||
self._target_publisher.publish(target_msg)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user