Part 6 Object Detection Node¶
Copy all the code below into your object_detection.py file, and make sure you read the annotations!
#!/usr/bin/env python3
import rclpy
from rclpy.node import Node # (1)!
import cv2
from cv_bridge import CvBridge, CvBridgeError # (2)!
from sensor_msgs.msg import Image # (3)!
from pathlib import Path # (4)!
class ObjectDetection(Node):
def __init__(self): # (5)!
super().__init__("object_detection")
self.camera_sub = self.create_subscription(
msg_type=Image,
topic="/camera/image_raw",
callback=self.camera_callback,
qos_profile=10
)
self.waiting_for_image = True # (6)!
def camera_callback(self, img_data): # (7)!
cvbridge_interface = CvBridge() # (8)!
try:
cv_img = cvbridge_interface.imgmsg_to_cv2(
img_data, desired_encoding="bgr8"
) # (9)!
except CvBridgeError as e:
self.get_logger().warning(f"{e}")
if self.waiting_for_image: # (10)!
height, width, channels = cv_img.shape
self.get_logger().info(
f"Obtained an image of height {height}px and width {width}px."
)
self.show_image(img=cv_img, img_name="step1_original")
self.waiting_for_image = False # (15)!
cv2.destroyAllWindows() # (16)!
def show_image(self, img, img_name, save_img=True): # (11)!
self.get_logger().info("Opening the image in a new window...")
cv2.imshow(img_name, img) # (12)!
if save_img: # (13)!
self.save_image(img, img_name)
self.get_logger().info(
"IMPORTANT: Close the image pop-up window to exit."
)
cv2.waitKey(0) # (14)!
def save_image(self, img, img_name): # (17)!
self.get_logger().info(f"Saving the image...")
base_image_path = Path.home().joinpath("myrosdata/object_detection/")
base_image_path.mkdir(parents=True, exist_ok=True) # (18)!
full_image_path = base_image_path.joinpath(
f"{img_name}.jpg") # (19)!
cv2.imwrite(str(full_image_path), img) # (20)!
self.get_logger().info(
f"\nSaved an image to '{full_image_path}'\n"
f" - image dims: {img.shape[0]}x{img.shape[1]}px\n"
f" - file size: {full_image_path.stat().st_size} bytes"
) # (21)!
def main(args=None):
rclpy.init(args=args)
node = ObjectDetection()
while node.waiting_for_image:
rclpy.spin_once(node) # (22)!
node.destroy_node()
rclpy.shutdown()
if __name__ == '__main__':
main()
-
Nothing new here, moving on...
-
We're importing the OpenCV library for Python (remember the Python API that we talked about earlier), which is called
cv2, and also that ROS-to-OpenCV bridge interface that we talked about earlier too:cv_bridge.From
cv_bridgewe're importing theCvBridgeandCvBridgeErrorclasses from thecv_bridgelibrary specifically. -
We need to subscribe to an image topic in order to obtain the data being published to it. You should've already identified the type of interface that is published to the
/camera/image_rawtopic, so we import that interface type here (from thesensor_msgspackage) so that we can build a subscriber to the topic later. -
We're also importing the Python
Pathclass from thepathlibmodule. A very handy tool for doing file operations. -
Initialising our
ObjectDetection()Class (should be very familiar to you by now):- Giving our node a name.
- Creating a subscriber to the
/camera/image_rawtopic, providing the interface type used by the topic (sensor_msgs/msg/Image- as imported above), and pointing it to a callback function (camera_callback, in this case), to define the processes that should be performed every time a message is obtained on this topic (in this case, the messages will be our camera images)
-
We're creating a flag to indicate whether the node has obtained an image yet or not. For this exercise, we only want to obtain a single image, so we will set the
waiting_for_imageflag toFalseonce an image has been obtained and processed, to avoid capturing any more.This flag will then be used to shut down the node when it's done its job.
-
Here, we're defining a callback function for our
self.camera_subsubscriber... -
Here, we create an instance of the
CvBridgeclass that we imported earlier, and which we'll use later on to convert ROS image data into a format that OpenCV can understand. -
We're using the CvBridge interface to take our ROS image data and convert it to a format that OpenCV will be able to understand. In this case we are specifying conversion (or "encoding") to an 8-bit BGR (Blue-Green-Red) image format:
"bgr8".We contain this within a
try-exceptblock though, which is the recommended procedure when doing this. Here we try to convert an image using the desired encoding, and if aCvBridgeErroris raised then we print a warning message to the terminal. Should this happen, this particular execution of the camera callback function will stop. -
Then we check the
waiting_for_imageflag to see if this is the first image that has been received by the node. If so, then:- Obtain the height and width of the image (in pixels), as well as the number of colour channels.
- Print a log message containing the image dimensions.
- Pass the image data to the
show_image()function (defined below). We also pass a descriptive name for the image to this function too (img_name).
-
This class method presents the image to us in a pop-up window and also calls another method which saves the image to file for us.
-
Display the actual image in a pop-up window:
- The image data is passed into the function via the
imgargument, - We need to give the pop-up window a name, so in this case we are using the
img_nameargument that is passed into this class method.
- The image data is passed into the function via the
-
The
show_image()class method has asave_imgargument, which is set toTrueby default, so that thisifcondition is triggered, and another class method is called to save the image to file. -
We're supplying a value of
0here, which tells this function to keep this window open indefinitely and wait until it is closed manually before allowing ourshow_image()class method to complete.If we had supplied a value here (say:
1) then the function would simply wait for 1 millisecond and then close the pop-up window down. In our case however, we want some time to actually look at the image and then close the window down ourselves, manually.Once the window has been closed, the execution of our code is able to continue...
-
We then set the
waiting_for_imageflag toFalseso that we only ever perform these processing steps once (we only want to capture a single image). This will then trigger the mainwhileloop to stop (see below), thus causing the overall execution of the node to stop too. -
cv2.destroyAllWindows()ensures that any OpenCV image pop-up windows that may still be active or in memory are destroyed before the class method exits (and the node shuts down). -
This class method handles the saving of the image to a file using
cv2tools andpathlib. -
Here, we define a filesystem location to save images to.
We want this to exist in a folder called "
myrosdata/object_detection" in the home directory, so we can use Pathlib'sPath.home().joinpath(...)to define it (a handy way to access the User's home directory, without needing to know the Users name).Then, we use the Pathlib
Path.mkdir()method to create this directory if it doesn't exist already. -
A full file path is constructed for the image here (using the
Path.joinpath()method), based on:- The
base_image_paththat we defined above - An image name that is passed into this class method via the
img_nameargument.
- The
-
This saves the image to a
.jpgfile. We're supplying thefull_image_paththat was created above, and also the actual image data (self.cv_img) so that the function knows what image we want to save. -
We're printing a log message to the terminal to inform us of:
- Where the image has been saved to
- How big the image is (in terms of its pixel dimensions)
- How big the image file is (in bytes).
-
We're using
spin_once()inside awhileloop here so that we can keep an eye on the value of thewait_for_imageflag, and stop spinning (i.e. break out of thewhileloop) once it turnsFalse.