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_bridge
we're importing theCvBridge
andCvBridgeError
classes from thecv_bridge
library 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_raw
topic, so we import that interface type here (from thesensor_msgs
package) so that we can build a subscriber to the topic later. -
We're also importing the Python
Path
class from thepathlib
module. 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_raw
topic, 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_image
flag toFalse
once 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_sub
subscriber... -
Here, we create an instance of the
CvBridge
class 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
-except
block though, which is the recommended procedure when doing this. Here we try to convert an image using the desired encoding, and if aCvBridgeError
is 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_image
flag 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
img
argument, - We need to give the pop-up window a name, so in this case we are using the
img_name
argument that is passed into this class method.
- The image data is passed into the function via the
-
The
show_image()
class method has asave_img
argument, which is set toTrue
by default, so that thisif
condition is triggered, and another class method is called to save the image to file. -
We're supplying a value of
0
here, 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_image
flag toFalse
so that we only ever perform these processing steps once (we only want to capture a single image). This will then trigger the mainwhile
loop 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
cv2
tools 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_path
that we defined above - An image name that is passed into this class method via the
img_name
argument.
- The
-
This saves the image to a
.jpg
file. We're supplying thefull_image_path
that 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 awhile
loop here so that we can keep an eye on the value of thewait_for_image
flag, and stop spinning (i.e. break out of thewhile
loop) once it turnsFalse
.