apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

adding synced data for pose estimation

Browse Source
This commit is contained in:
Apoorva Gupta 2023-03-22 11:20:49 +05:30
parent cf0fdc175b
commit 61e7117511
1 changed files with 36 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
ros2_ws/src/yolov3/detect_on_bag.py
View File

@ -43,6 +43,7 @@ from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterb
import numpy as np import numpy as np
from yolov3_msg.msg import BoundingBox, BoundingBoxes from yolov3_msg.msg import BoundingBox, BoundingBoxes
from rclpy.clock import Clock from rclpy.clock import Clock
from sensor_msgs.msg import CameraInfo
@torch.no_grad() @torch.no_grad()
def run(weights=ROOT / 'yolov3.pt', # model.pt path(s) def run(weights=ROOT / 'yolov3.pt', # model.pt path(s)
@ -161,27 +162,59 @@ class DetectOnBag(Node):
def __init__(self): def __init__(self):
super().__init__('detect_on_bag') super().__init__('detect_on_bag')
self.camera_info = None
self.depth_img = None
self.subscription = self.create_subscription( self.subscription = self.create_subscription(
Image, Image,
'/camera/color/image_raw', '/camera/color/image_raw',
self.image_callback, self.image_callback,
10) 10)
self.subscription # prevent unused variable warning self.subscription # prevent unused variable warning
self.subscription = self.create_subscription(
Image,
'/camera/aligned_depth_to_color/image_raw',
self.image_depth_callback,
10)
self.subscription
self.subscription = self.create_subscription(
CameraInfo,
'/camera/color/camera_info',
self.camera_info_callback,
10)
self.subscription
self.detected_img_pub = self.create_publisher(Image, '/detection_image', 10) self.detected_img_pub = self.create_publisher(Image, '/detection_image', 10)
self.bboxes_pub = self.create_publisher(BoundingBoxes, '/bboxes', 10) self.bboxes_pub = self.create_publisher(BoundingBoxes, '/bboxes', 10)
self.depth_pub = self.create_publisher(Image, '/depth_img', 10)
self.img_pub = self.create_publisher(Image, '/rgb_img', 10)
self.info_pub = self.create_publisher(CameraInfo, '/camera_info', 10)
def camera_info_callback(self, msg):
print("camera info cb")
# print(msg.header.stamp)
self.camera_info = msg
def image_depth_callback(self, msg):
print("depth info cb")
# print(msg.header.stamp)
self.depth_img = msg
def image_callback(self, msg): def image_callback(self, msg):
self.get_logger().info('Image') print("rgb img cb")
print(msg.header.stamp) # print(msg.header.stamp)
bridge = CvBridge() bridge = CvBridge()
cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough') cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
detected_img, bboxes = run('src/yolov3/runs/train/exp14/weights/best.pt', cv_image) detected_img, bboxes = run('src/yolov3/runs/train/exp15/weights/best.pt', cv_image)
self.detected_img_pub.publish(self.numpy_array_to_image_msg(detected_img)) self.detected_img_pub.publish(self.numpy_array_to_image_msg(detected_img))
bboxes.image_header = msg.header bboxes.image_header = msg.header
bboxes.header.stamp = msg.header.stamp bboxes.header.stamp = msg.header.stamp
bboxes.header.frame_id = 'detection' bboxes.header.frame_id = 'detection'
self.bboxes_pub.publish(bboxes) self.bboxes_pub.publish(bboxes)
if self.depth_img:
self.depth_pub.publish(self.depth_img)
self.img_pub.publish(msg)
if self.camera_info:
self.info_pub.publish(self.camera_info)
def numpy_array_to_image_msg(self, numpy_array): def numpy_array_to_image_msg(self, numpy_array):
# Create a CvBridge object # Create a CvBridge object