예제 #1
0
def handle_recognize(req):
    rsp = recognizeResponse()

    bridge = CvBridge()
    try:
        # Convert rosmsg to cv image
        # np_array = np.fromstring(req.img_in.data, np.uint8)
        # image = cv2.imdecode(np_array, cv2.CV_LOAD_IMAGE_COLOR)

        cv_image = bridge.imgmsg_to_cv2(req.img_in, "bgr8")
        (rows, cols, channels) = cv_image.shape

        if cols != 640 or rows != 480:
            rospy.WARN('Can not get image.')
            return rsp
        else:
            result = ps.process(cv_image)
            ro_msg = recognized_objects()

            for i in result:
                bbox = i[0]
                score = i[1]
                class_name = i[2]

                cv2.rectangle(cv_image, (int(bbox[0]), int(bbox[1])),
                              (int(bbox[2]), int(bbox[3])), (255, 0, 255), 2)
                cv2.putText(cv_image, '{:s} {:.3f}'.format(class_name, score),
                            (int(bbox[0]), int(bbox[1]) - 2),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

                label_str_msg = String()
                label_str_msg.data = class_name

                bbox_array_msg = UInt16MultiArray()
                for j in bbox:
                    bbox_array_msg.data.append(j)

                ro_msg.labels.append(label_str_msg)
                ro_msg.bbox_arrays.append(bbox_array_msg)

            ro_msg.header = req.img_in.header

            rsp.img_out = bridge.cv2_to_imgmsg(cv_image, "bgr8")
            rsp.obj_info = ro_msg
            return rsp

    except CvBridgeError as e:
        print(e)
        return rsp
    def _callback(self, req):
        if self._init is False:
            self._net = load_net()
            self._init = True

        use_gpu = True
        if use_gpu:
            caffe.set_mode_gpu()
            caffe.set_device(0)
            cfg.GPU_ID = 0
        else:
            caffe.set_mode_cpu()

        # Convert rosmsg to cv image
        bridge = CvBridge()
        image = bridge.imgmsg_to_cv2(req.img_in, "bgr8")

        objects = self._detect(image)

        ro_msg = msg.recognized_objects()
        rsp = srv.recognizeResponse()

        for i in objects:
            bbox = i[0]
            class_name = i[2]

            label_str_msg = String()
            label_str_msg.data = class_name

            bbox_array_msg = UInt16MultiArray()
            for j in bbox:
                bbox_array_msg.data.append(j)

            ro_msg.labels.append(label_str_msg)
            ro_msg.bbox_arrays.append(bbox_array_msg)

        ro_msg.header = req.img_in.header

        rsp.obj_info = ro_msg
        return rsp