로봇/ROS

ROS1에 YOLO v5 연동하기

with-RL 2023. 1. 18. 20:44

카메라는 로봇이 주변 상황을 인식하기 위해 사용되는 장비로 사람의 눈과 같은 역할을 합니다. 카메라로 인식된 영상을 YOLO v5를 이용해 Object Detection을 실행하고 그 결과를 publish 하고 subscribe 과정에 대한 설명입니다.

 

아래 과정은 모두 아래 명령을 통해 Docker Container에 접속한 이후 실행하는 내용입니다.

$ docker exec -it ros1 /bin/bash

1. 시험환경

2. ROS1에 카메라 연동하기

2.1. roscore 실행

아래 명령을 실행해서 roscore를 실행합니다.

# roscore

2.2. uvc-camera 연동

roscore가 실행 중인 터미널이 아닌 새로운 터미널에서 아래명령을 통해 'uvc-camera' node를 설치합니다.

# apt install ros-melodic-uvc-camera

아래 명령을 통해 uvc-camera-node를 실행합니다.

아래 명령 중 '_skip_frames:=5'는 파라미터 값으로 6 프레임 중 5 프레임을 쉬고 1 프레임을 전송하라는 의미입니다.

이 명령을 준 이유는 노트북에 부하를 줄이기 위함이고 기본값은 '_skip_frames:=0'입니다.

# rosrun uvc_camera uvc_camera_node _skip_frames:=5

아래 명령을 통해서 카메라의 영상을 확인할 수 있습니다.

# rosrun image_view image_view image:=image_raw

실행된 카메라 영상은 초당 30 프레임 중 5 프레임만 보낸 것으로 영상이 끊어져서 보이면 정상입니다.

아래 명령을 통해서 ROS1의 topic의 흐름을 확인할 수 있습니다.

# rqt_graph

위 그림과 같이 '/uvc_camera' node에서 '/image_view' node로 '/image_raw' topic을 전송하고 있음을 확인할 수 있습니다.

2.3. ROS Workspace 생성

ROS 프로그래밍을 위해서 Ubuntu에서 Docker를 이용한 ROS1, ROS2 설치 과정에서 지정한 <Container_ws>로 이동합니다.

저는 <Container_ws>를 '/home/cchyun/Workspace'로 지정했습니다. 본인의 환경에 맞게 수정해서 사용하시면 됩니다.

# cd /home/cchyun/Workspace

다음 명령을 통해 ros_ws와 그 아래 src 폴더를 생성하고 ros_ws 폴더로 이동합니다.

# mkdir -p ros_ws/src
# cd ros_ws

ROS 컴파일 명령어인 catkin_make 함수를 실행해서 컴파일합니다.

# catkin_make

아직 추가된 코드는 없지만 기본적인 빌드가 됩니다.

2.4. yolov5_ros 패키지 기능 구현

아래 명령을 실행해서 'yolov5_ros' 패키지를 생성합니다. 이 패키지는 uvc_camera node에서 받은 이미지를 YOLO v5로 Object Detection을 수행하고 결과를 publish 하는 역할을 수행합니다.

# cd /home/cchyun/Workspace/ros_ws/src
# catkin_create_pkg yolov5_ros rospy roscpp std_msgs sensor_msgs message_generation

ls 명령을 통해서 확인해 보면 yolov5_ros 폴더가 생성된 것을 확인할 수 있습니다.

이제 아래 명령을 통해 yolov5를 다운로드합니다.

# cd home/cchyun/Workspace/ros_ws/src/yolov5_ros/src
# git clone https://github.com/ultralytics/yolov5

아래 명령을 통해 yolov5를 설치합니다.

# cd yolov5
# pip3 install --ignore-installed PyYAML
# pip3 install -r requirements.txt

이제 yolov5_ros에서 Object Detection 결과를 전송 하기 위한 msg를 정의하기 위해 msg 폴더를 생성합니다.

# cd /home/cchyun/Workspace/ros_ws/src/yolov5_ros
# mkdir msg

https://github.com/raghavauppuluri13/yolov5_pytorch_ros/tree/master/msg에서 BoundingBox.msg, BoundingBoxes.msg 두 파일을 다운로드하여 msg 폴더에 저장합니다. 이 파일에는 ROS를 통해서 전송할 Topic 메시지 규격이 정리돼 있습니다.

# wget https://raw.githubusercontent.com/raghavauppuluri13/yolov5_pytorch_ros/master/msg/BoundingBox.msg
# wget https://raw.githubusercontent.com/raghavauppuluri13/yolov5_pytorch_ros/master/msg/BoundingBoxes.msg

다음은 아래 내용을 '/home/cchyun/Workspace/ros_ws/src/yolov5_ros/CMakeLists.txt' 파일의 'Declare ROS messages, services and actions' 부분 아래쪽에 기록합니다.

################################################
## Declare ROS messages, services and actions ##
################################################
# ...

add_message_files(
  FILES
  BoundingBox.msg
  BoundingBoxes.msg
)

generate_messages(
  DEPENDENCIES
    geometry_msgs
    sensor_msgs
    std_msgs
)

다음은 catkin_make를 실행합니다.

# cd /home/cchyun/Workspace/ros_ws
# catkin_make

다음은 첨부된 파일을 '/home/cchyun/Workspace/ros_ws/src/yolov5_ros/src'에 detect.py 파일을 아래 내용으로 생성합니다.

소스코드 https://github.com/mats-robotics/yolov5_ros/blob/main/src/detect.py을 약간 수정했습니다.

#!/usr/bin/env python3
# 참조: https://github.com/mats-robotics/yolov5_ros/blob/main/src/detect.py
# 아래 내용은 위 파일의 내용을 기반으로 필요한 부분만을 최소한으로 작성한 파일 입니다.

import sys
import os
import time
import rospy
import numpy as np
import torch
import torch.backends.cudnn as cudnn

from pathlib import Path

from sensor_msgs.msg import Image
from yolov5_ros.msg import BoundingBox, BoundingBoxes

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] / "yolov5"
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative path

from models.common import DetectMultiBackend
from utils.general import (
    check_img_size,
    non_max_suppression,
    scale_boxes
)
from utils.torch_utils import select_device
from utils.augmentations import letterbox


class Yolov5Detector:
    def __init__(self):
        self.conf_thres = rospy.get_param("~confidence_threshold", 0.75)
        self.iou_thres = rospy.get_param("~iou_threshold", 0.45)
        self.agnostic_nms = rospy.get_param("~agnostic_nms", True)
        self.max_det = rospy.get_param("~maximum_detections", 1000)
        self.classes = rospy.get_param("~classes", None)
        # Initialize weights 
        weights = rospy.get_param("~weights", "yolov5s.pt")
        # Initialize model
        self.device = select_device(str(rospy.get_param("~device","cpu")))
        self.model = DetectMultiBackend(weights, device=self.device, dnn=rospy.get_param("~dnn", True), data=rospy.get_param("~data", ""))
        self.stride, self.names, self.pt, self.jit, self.onnx, self.engine = (
            self.model.stride,
            self.model.names,
            self.model.pt,
            self.model.jit,
            self.model.onnx,
            self.model.engine,
        )

        # Setting inference size
        self.img_size = [rospy.get_param("~inference_size_w", 640), rospy.get_param("~inference_size_h", 480)]
        self.img_size = check_img_size(self.img_size, s=self.stride)

        # Half
        self.half = rospy.get_param("~half", False)
        self.half &= (
            self.pt or self.jit or self.onnx or self.engine
        ) and self.device.type != "cpu"  # FP16 supported on limited backends with CUDA
        if self.pt or self.jit:
            self.model.model.half() if self.half else self.model.model.float()
        bs = 1  # batch_size
        cudnn.benchmark = True  # set True to speed up constant image size inference
        self.model.warmup(imgsz=(1 if self.pt else bs, 3, *self.img_size))  # warmup        
        
        # Initialize subscriber to Image topic
        rospy.Subscriber('/image_raw', Image, self.callback)
        self.pred_pub = rospy.Publisher(rospy.get_param("~output_topic", "/yolov5_ros"), BoundingBoxes, queue_size=10)

    def callback(self, img_msg):
        dtype = np.dtype("uint8")
        dtype = dtype.newbyteorder('>' if img_msg.is_bigendian else '<')
        img_cv = np.ndarray(shape=(img_msg.height, img_msg.width, 3), dtype=dtype, buffer=img_msg.data)
        if img_msg.is_bigendian == (sys.byteorder == 'little'):
            img_cv = img_cv.byteswap().newbyteorder()
        im, im0 = self.preprocess(img_cv)

        im = torch.from_numpy(im).to(self.device) 
        im = im.half() if self.half else im.float()
        im /= 255
        if len(im.shape) == 3:
            im = im[None]
    
        pred = self.model(im, augment=False, visualize=False)
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, self.agnostic_nms, max_det=self.max_det)

        det = pred[0].cpu().numpy()

        bounding_boxes = BoundingBoxes()
        bounding_boxes.header = img_msg.header
        bounding_boxes.image_header = img_msg.header

        if len(det):
            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
            for *xyxy, conf, cls in reversed(det):
                c = int(cls)

                bbox_msg = BoundingBox()
                bbox_msg.Class = self.names[c]
                bbox_msg.probability = conf 
                bbox_msg.xmin = int(xyxy[0])
                bbox_msg.ymin = int(xyxy[1])
                bbox_msg.xmax = int(xyxy[2])
                bbox_msg.ymax = int(xyxy[3])

                bounding_boxes.bounding_boxes.append(bbox_msg)
        
        self.pred_pub.publish(bounding_boxes)
    
    def preprocess(self, img):
        img0 = img.copy()
        img = np.array([letterbox(img, self.img_size, stride=self.stride, auto=self.pt)[0]])
        img = img[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
        img = np.ascontiguousarray(img)

        return img, img0 


if __name__ == "__main__":
    rospy.init_node("yolov5", anonymous=True)

    detector = Yolov5Detector()
    
    rospy.spin()

이제 detect.py 파일에 실행 권한을 줍니다.

# cd /home/cchyun/Workspace/ros_ws/src/yolov5_ros/src
# chmod +x detect.py

다음은 아래 내용을 '/home/cchyun/Workspace/ros_ws/src/yolov5_ros/CMakeLists.txt' 파일의 'Install' 아랫부분에 추가합니다.

#############
## Install ##
#############
# ...
catkin_install_python(PROGRAMS
  src/detect.py
  DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
)

이제 catkin_make를 실행하고 detect.py를 실행합니다.

# cd /home/cchyun/Workspace/ros_ws
# catkin_make
# source devel/setup.bash
# rosrun yolov5_ros detect.py

rqt_graph를 다시 실행해서 Topic의 흐름을 확인합니다.

# rqt_graph

'/uvc_camera' node에서 발생한 Topic을 '/Image_view' node와 '/yolov5_ros' node에서 수신하는 것을 확인할 수 있습니다.

2.5. robot_ctrl 패키지 기능 구현

다음은 '/uvc_camera' node와 '/yolov5_ros' node의 정보를 받아서 로봇을 제어하기 위한 패키지를 만듭니다. 아직 구체적인 제어할 로봇이 결정되지 않았으므로 수신한 정보를 확인하는 기능까지만 구현하겠습니다.

우선 아래 명령을 통해 robot_ctrl 패키지를 생성합니다.

# cd /home/cchyun/Workspace/ros_ws/src
# catkin_create_pkg robot_ctrl rospy roscpp std_msgs sensor_msgs yolov5_ros

'/home/cchyun/Workspace/ros_ws/src/robot_ctrl/src' 폴더아래 control.py를 생성하고 아래 내용을 붙여 넣습니다.

#!/usr/bin/env python3

import sys
import numpy as np
import cv2
import rospy
import message_filters

from sensor_msgs.msg import Image
from yolov5_ros.msg import BoundingBoxes

class RobotControl:
    def __init__(self):
        image_cb = message_filters.Subscriber('/image_raw', Image)
        yolov5_cb = message_filters.Subscriber('/yolov5_ros', BoundingBoxes)

        ts = message_filters.TimeSynchronizer([image_cb, yolov5_cb], 10)
        ts.registerCallback(self.callback)

    def callback(self, img_msg, bbox_msg):
        dtype = np.dtype("uint8")
        dtype = dtype.newbyteorder('>' if img_msg.is_bigendian else '<')
        img_cv = np.ndarray(shape=(img_msg.height, img_msg.width, 3), dtype=dtype, buffer=img_msg.data)
        if img_msg.is_bigendian == (sys.byteorder == 'little'):
            img_cv = img_cv.byteswap().newbyteorder()
        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
        for bbox in bbox_msg.bounding_boxes:
            print(bbox)
            cv2.rectangle(img_cv, (bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax), (255,0,0), 2)

        cv2.imwrite('aaa.jpg', img_cv)


if __name__ == "__main__":
    rospy.init_node("robot_control", anonymous=True)

    control = RobotControl()
    
    rospy.spin()

control.py에 실행 권한을 줍니다.

# cd /home/cchyun/Workspace/ros_ws/src/robot_ctrl/src
# chmod +x control.py

다음은 아래 내용을 '/home/cchyun/Workspace/ros_ws/src/robot_ctrl/CMakeLists.txt' 파일의 'Install' 아랫부분에 추가합니다.

 

#############
## Install ##
#############
# ...
catkin_install_python(PROGRAMS
  src/control.py
  DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
)

이제 catkin_make를 실행하고 control.py를 실행합니다.

# cd /home/cchyun/Workspace/ros_ws
# catkin_make
# source devel/setup.bash
# rosrun robot_ctrl control.py

rqt_graph를 다시 실행해서 Topic의 흐름을 확인합니다.

# rqt_graph

'/robot_control' node는 '/uvc_camera' node와 '/yolov5_ros' node에서 topic을 받은 것을 확인할 수 있습니다.

'/home/cchyun/Workspace/ros_ws/aaa.jpg' 파일을 확인하면 Object Detection 된 결과를 확인할 수 있습니다.

이상으로 ROS1에 yolov5를 연동하는 과정에 대한 설명이었습니다.