runtime/rbss_objectdetection/scripts/od_yolo_lc.py

#!/usr/bin/env python3
"""
  object_detection_lifecycle_node_via_interface_node
  ROS 2 program for Object Detection

  @shalenikol release 0.4
  ...
"""

from typing import Optional
import os
import json

import rclpy
from rclpy.lifecycle import Node
from rclpy.lifecycle import Publisher
from rclpy.lifecycle import LifecycleState
from rclpy.lifecycle import TransitionCallbackReturn
from lifecycle_msgs.msg import State
# from lifecycle_msgs.srv import GetState
from rclpy.timer import Timer

from ament_index_python.packages import get_package_share_directory
from sensor_msgs.msg import Image
from rbs_skill_interfaces.msg import BoundBox
from rbs_skill_interfaces.srv import RbsBt

from cv_bridge import CvBridge # Package to convert between ROS and OpenCV Images

# CAMERA_LINK_DEFAULT = "rgbd_camera"
CONDITION_SRV_NAME = "/condition"
FILE_DOPE_CONFIG = "od_yolo_config.yaml"
NODE_NAME_DEFAULT = "lc_yolo"
PARAM_SKILL_CFG = "lc_yolo_cfg"
OUT_TOPIC_DEFAULT = "/object_detection"
OUT_TOPIC_IMAGE = "/detect_image"
OUT_TOPIC_TYPE = "rbs_skill_interfaces/msg/BoundBox"

class ObjectDetection(Node):
    """Our lifecycle node."""
    def __init__(self, **kwargs):
        """Construct the node."""
        self._count: int = 0
        self._pub: Optional[Publisher] = None
        self._pubI: Optional[Publisher] = None
        self._timer: Optional[Timer] = None
        self._image_cnt: int = 0
        self._sub = None
        self._threshold = 0.0
        self._is_image_mode = False
        self.yolov8_weights_file = ""
        self.model = None
        self.is_detection = False # object detection indicator in the frame

        # for other nodes
        kwargs["allow_undeclared_parameters"] = True
        kwargs["automatically_declare_parameters_from_overrides"] = True
        super().__init__(NODE_NAME_DEFAULT, **kwargs)
        str_cfg = self.get_parameter(PARAM_SKILL_CFG).get_parameter_value().string_value

        self.skill_cfg = json.loads(str_cfg)
        self.nodeName = NODE_NAME_DEFAULT # self.skill_cfg["Module"]["node_name"] #["Launch"]["name"]

        self.topicImage = "" # input topic
        self.publishDelay = 0.33
        self.topicSrv = self.nodeName + OUT_TOPIC_DEFAULT # self.skill_cfg["topicsOut"][0]["name"] # self.nodeName + "/" + self.skill_cfg["Interface"]["Output"][0]["name"]
        self.topicDetectImage = self.nodeName + OUT_TOPIC_IMAGE
        # self._Settings()

        # Used to convert between ROS and OpenCV images
        self.br = CvBridge()

        self.objName = ""
        self.image_det = []
        self.bbox_res = None

        self._srv_condition = self.create_service(RbsBt, NODE_NAME_DEFAULT + CONDITION_SRV_NAME, self.condition_callback)
        # , callback_group=self.cb_group)

    def _Settings(self):
        # Initialization service settings
        for prop in self.skill_cfg["Settings"]["output"]["params"]:
            nam = prop["name"]
            val = prop["value"]
            if nam == "threshold":
                self._threshold = float(val)
            elif nam == "publishDelay":
                self.publishDelay = float(val)
            elif nam == "is_image_mode":
                self._is_image_mode = (val == "True")

        for prop in self.skill_cfg["topicsOut"]:
            if prop["type"] == OUT_TOPIC_TYPE:
                self.topicSrv = prop["name"]

    def condition_callback(self, request, response):
        _is = False
        if request.command == "isDetectionRun":
            # self._cli_getstate = self.create_client(GetState, f"{self.nodeName}/get_state")
            # req = GetState.Request()
            # self.get_logger().info(f"condition_callback : before call")
            # res = self._cli_getstate.call(req) # the program freezes
            id_,_ = self._state_machine.current_state
            _is = (id_ == State.PRIMARY_STATE_ACTIVE)
        elif request.command == "isDetection":
            _is = self.is_detection
        response.ok = _is
        return response

    def on_configure(self, state: LifecycleState) -> TransitionCallbackReturn:
        """
        Configure the node, after a configuring transition is requested.

        :return: The state machine either invokes a transition to the "inactive" state or stays
            in "unconfigured" depending on the return value.
            TransitionCallbackReturn.SUCCESS transitions to "inactive".
            TransitionCallbackReturn.FAILURE transitions to "unconfigured".
            TransitionCallbackReturn.ERROR or any uncaught exceptions to "errorprocessing"
        """
        json_param = self.get_parameter(PARAM_SKILL_CFG).get_parameter_value().string_value
        self.skill_cfg = json.loads(json_param)
        self._Settings()
        dependency = {}
        for comm in self.skill_cfg["BTAction"]:
            for par in comm["param"]:
                dependency = par["dependency"]
                if par["type"] == "topic":
                    self.topicImage = dependency["topicOut"]
                elif par["type"] == "formBuilder":
                    self.objName = dependency["output"]["object_name"]
                    val = dependency["output"]["process"]["selectProcess"]["value"]
                    self.yolov8_weights_file = os.path.join(val["instancePath"], val["instanceName"]+".pt")
        assert self.objName, "no object name"
        assert self.topicImage, "no input topic dependency"
        # assert self.yolov8_weights_file, "no weights dependency"

        # self.yolov8_weights_file = dependency["weights_file"]
        if not os.path.isfile(self.yolov8_weights_file):
            self.get_logger().warning(f"No weights found <{self.yolov8_weights_file}>")
            return TransitionCallbackReturn.FAILURE

        # Create the publisher.
        self._pub = self.create_lifecycle_publisher(BoundBox, self.topicSrv, 10)
        if self._is_image_mode:
            self._pubI = self.create_lifecycle_publisher(Image, self.topicDetectImage, 1)

        self._timer = self.create_timer(self.publishDelay, self.publish)

        from ultralytics import YOLO
        self.model = YOLO(self.yolov8_weights_file)

        self.get_logger().info("on_configure() is called.")
        return TransitionCallbackReturn.SUCCESS

    def on_activate(self, state: LifecycleState) -> TransitionCallbackReturn:
        self.get_logger().info('on_activate() is called.')
        # Create the main subscriber.
        self._sub = self.create_subscription(Image, self.topicImage, self.listener_callback, 1)
        return super().on_activate(state)

    def on_deactivate(self, state: LifecycleState) -> TransitionCallbackReturn:
        self.get_logger().info('on_deactivate() is called.')
        # Destroy the main subscriber.
        self.destroy_subscription(self._sub)
        return super().on_deactivate(state)

    def on_cleanup(self, state: LifecycleState) -> TransitionCallbackReturn:
        """
        Cleanup the node.
        """
        # # очистим параметры
        # node_param = rclpy.parameter.Parameter("setting_path", rclpy.Parameter.Type.STRING, "")
        # all_node_param = [node_param]
        # self.set_parameters(all_node_param)

        self._is_image_mode = False
        self.image_det = []
        self.bbox_res = None
        self.topicImage = ""
        self.objName = ""
        self.yolov8_weights_file = ""

        self.destroy_timer(self._timer)
        self.destroy_publisher(self._pub)
        self.destroy_publisher(self._pubI)
        self.destroy_subscription(self._sub)

        self.get_logger().info('on_cleanup() is called.')
        return TransitionCallbackReturn.SUCCESS

    def on_shutdown(self, state: LifecycleState) -> TransitionCallbackReturn:
        """
        Shutdown the node.
        """
        self.destroy_timer(self._timer)
        self.destroy_publisher(self._pub)
        self.destroy_publisher(self._pubI)
        self.destroy_subscription(self._sub)

        self.get_logger().info('on_shutdown() is called.')
        return TransitionCallbackReturn.SUCCESS

    def publish(self):
        """Publish a new message when enabled."""
        self._count += 1

        if self._pub is not None and self._pub.is_activated:
            # Publish result
            if not self.bbox_res:
                return
            msg = BoundBox()
            msg.is_detection = False
            self.is_detection = False
            #from ultralytics.engine.results
            cconf = self._threshold # threshold
            bb = None
            for c in self.bbox_res.boxes:
                idx = int(c.cls)
                if self.bbox_res.names[idx] == self.objName:
                    conf = float(c.conf)
                    if cconf < conf:
                        cconf = conf
                        bb = c
            if bb:
                msg.is_detection = True
                self.is_detection = True
                msg.name = self.objName
                msg.x = float(bb.xywhn[0,0])
                msg.y = float(bb.xywhn[0,1])
                msg.w = float(bb.xywhn[0,2])
                msg.h = float(bb.xywhn[0,3])
                msg.conf = float(bb.conf)
             # Publish the message.
            self._pub.publish(msg)

            if self._is_image_mode:
                if len(self.image_det) == 0:
                    return
                # The 'cv2_to_imgmsg' method converts an OpenCV image to a ROS 2 image message
                msg = self.br.cv2_to_imgmsg(self.image_det, encoding="rgb8")
                # Publish the message.
                self._pubI.publish(msg)

    def listener_callback(self, data):
        """
        Image Callback function.
        """
        if self.model:
            self._image_cnt += 1
            if self._image_cnt % 200 == 1:
              self.get_logger().info(f"detection: {self._image_cnt}")

            # Convert ROS Image message to OpenCV image
            current_frame = self.br.imgmsg_to_cv2(data)
            # run inference
            results = self.model(current_frame)

            for r in results:
                self.bbox_res = r

            if self._is_image_mode:
                for r in results:
                    self.image_det = r.plot()  # plot a BGR numpy array of predictions

            # self.get_logger().info(f"detection: end {self._image_cnt}")

def main():
    rclpy.init()

    executor = rclpy.executors.SingleThreadedExecutor()
    lc_node = ObjectDetection()
    executor.add_node(lc_node)
    try:
        executor.spin()
    except (KeyboardInterrupt, rclpy.executors.ExternalShutdownException):
        lc_node.destroy_node()

if __name__ == '__main__':
    main()