Welcome to DOFBOT-Ultra repository

Model Prediction

Note: The environment is already configured. Simply follow the tutorial to run the corresponding functional commands.

1. Program Description

Use Python to call a USB camera or depth camera with your own trained model. Here we use our trained model for recognizing oranges as an example. If you need to train your own model, refer to the two links below, but training on the RDK X5 is not recommended.

Annotation: GitHub - HumanSignal/label-studio: Label Studio is a multi-type data labeling and annotation tool with standardized output format

Training: Model Training with Ultralytics YOLO - Ultralytics YOLO Docs

2. Program Startup

3.1. Start the Camera

Start the following program according to your camera model, input in the terminal:

xxxxxxxxxx
# USB camera
ros2 launch dofbot_ultra_info usb_cam.launch.py
# Depth camera
ros2 launch orbbec_camera dabai_dcw2.launch.py

Open another terminal and enter the code folder:

xxxxxxxxxx
cd /home/sunrise/ultralytics/ultralytics/yahboom_demo

Run the code. Click on the preview window and press the q key to terminate the program!

xxxxxxxxxx
python3 06.orange_camera_usb.py

Result Preview

YOLO recognition output video location: /home/sunrise/ultralytics/ultralytics/output/

Example code:

​x
#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CompressedImage
from cv_bridge import CvBridge
import cv2
from ultralytics import YOLO
import os
​
class Image_detection(Node):
    def __init__(self):
        super().__init__('Image_detection')
        self.model = YOLO("/home/sunrise/ultralytics/ultralytics/data/yahboom_data/orange_data/runs/detect/train/weights/best.onnx") # The path here is the path to your personal trained model
        self.camera_type = os.getenv('CAMERA_TYPE', 'usb')
        self.bridge = CvBridge()
        if self.camera_type == 'usb':
            topic_name = '/camera/color/image_raw' 
        else:
            topic_name = '/camera/color/image_raw'
​
        self.subscription = self.create_subscription(Image,topic_name, self.image_callback,10)
​
        # Get the video frame size and frame rate
        frame_width = 640
        frame_height = 480
        fps = 15
        
        output_path = "/home/sunrise/ultralytics/ultralytics/output/06.orange_camera_usb.mp4"
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can use 'XVID' or 'mp4v' depending on your platform
        self.out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
​
    def image_callback(self, msg):
        cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
 
        self.proecc(cv_image)
​
# Loop through the video frames
    def proecc(self,frame):
        # Run YOLO inference on the frame
        results = self.model(frame)
​
        # Visualize the results on the frame
        annotated_frame = results[0].plot()
​
        # Write the annotated frame to the output video file
        self.out.write(annotated_frame)
​
​
        # Display the annotated frame
        cv2.imshow("YOLO Inference", cv2.resize(annotated_frame, (640, 480)))
​
        # Break the loop if 'q' is pressed
        cv2.waitKey(1) & 0xFF == ord("q")
​
    def cancel(self):
        cv2.destroyAllWindows()
        self.out.release()
​
def main(args=None):
    rclpy.init(args=args)
    node = Image_detection()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.cancel()
        node.destroy_node()
        rclpy.shutdown()
​
if __name__ == '__main__':
    main()