drone-rigide/workspace/src/detect_targets/scripts/target_publisher.py

#! /usr/bin/python

import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np
import dynamic_reconfigure.server

from vqimg.msg import component_centers, component_center
from find_targets import find_targets, normalize_coordinates
from detect_targets.cfg import DetectTargetsConfig


class Publisher:
    def __init__(self):
        self.threshold_blue = 140
        self.threshold_red = 120
        self.threshold_green = 190
        self.centers = component_centers()
        self.pub = rospy.Publisher('targets', component_centers, queue_size=10)
        self.img_pub = rospy.Publisher('img_targets', Image, queue_size=0)
        rospy.Subscriber("/bebop/image_raw", Image, self.on_picture)
        self.bridge = CvBridge()
        self.config_srv = dynamic_reconfigure.server.Server(DetectTargetsConfig, self.on_reconf)
        self.draw_targets = False

    def on_reconf(self, config, level):
        self.threshold_red = config['red']
        self.threshold_blue = config['blue']
        self.threshold_green = config['green']
        self.draw_targets = config['draw_targets']
        return config

    def on_picture(self, msg):
        img = np.array(np.asarray(self.bridge.imgmsg_to_cv2(msg)), copy=True)
        try:
            H, L, R = find_targets(
                img,
                threshold_blue=self.threshold_blue,
                threshold_red=self.threshold_red,
                threshold_green=self.threshold_green,
            )
        except ValueError:
            return
        pub_H = normalize_coordinates(H, msg.width, msg.height)
        pub_L = normalize_coordinates(L, msg.width, msg.height)
        pub_R = normalize_coordinates(R, msg.width, msg.height)
        self.centers.data = [
            component_center(x=pub_H[0], y=pub_H[1], nb_vertex=1, label=1),
            component_center(x=pub_L[0], y=pub_L[1], nb_vertex=1, label=2),
            component_center(x=pub_R[0], y=pub_R[1], nb_vertex=1, label=3),
        ]

        if self.draw_targets:
            for i in range(-20, 20):
                for j in range(-20, 20):
                    try:
                        img[int(H[1])-i, int(H[0])-j] = [0, 255, 0]
                        img[int(L[1])-i, int(L[0])-j] = [255, 0, 0]
                        img[int(R[1])-i, int(R[0])-j] = [0, 0, 255]
                    except Exception:
                        pass

        self.pub.publish(self.centers)
        self.img_pub.publish(self.bridge.cv2_to_imgmsg(img, "rgb8"))


if __name__ == '__main__':
    rospy.init_node('target_publisher')

    publisher = Publisher()
    rospy.spin()