Welcome to ROSMASTER A1 repository

Model prediction

Note: Using the Docker container in the factory image does not require re-setting up the environment. The environment is already set up. Simply follow the previous tutorial to enter Docker and run the corresponding function commands.

1. Program Description

Use Python to call a USB camera or depth camera using a pre-trained module. Here, we use the model we trained to recognize oranges as an example. If you need to train your own model, please refer to the two links below. However, training on a Raspberry Pi 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. Camera Startup

Start the following program based on your camera model. In the terminal, enter:

xxxxxxxxxx
#usb camera
ros2 launch usb_cam camera.launch.py
#nuwa camera
ros2 launch ascamera hp60c.launch.py

Open another terminal and navigate to the code folder:

xxxxxxxxxx
cd /root/ultralytics/ultralytics/yahboom_demo

Run the code: Click the preview screen and press q to terminate the program!

xxxxxxxxxx
python3 06.orange_camera_usb.py

Preview

Yolo recognition output video location: /root/ultralytics/ultralytics/output/

Sample code:

xxxxxxxxxx
#!/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("/root/ultralytics/ultralytics/data/yahboom_data/orange_data/runs/detect/train/weights/best_ncnn_model") #这里的路径是个人训练模型的路径 The path here is the path of the personal training model
        self.camera_type = os.getenv('CAMERA_TYPE', 'usb')
        self.bridge = CvBridge()
        if self.camera_type == 'usb':
            topic_name = '/usb_cam/image_raw' 
        else:
            topic_name = '/ascamera_hp60c/camera_publisher/rgb0/image' 
​
        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 = "/root/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()