src/jsk_perception/label_image_classifier.py

Go to the documentation of this file.

#!/usr/bin/env python
 
import cv_bridge
from jsk_recognition_msgs.msg import ClassificationResult
from jsk_topic_tools import ConnectionBasedTransport
import numpy as np
import rospy
from sensor_msgs.msg import Image
 
class LabelImageClassifier(ConnectionBasedTransport):
 
    classifier_name = 'label_image_classifier'
 
    def __init__(self):
        super(LabelImageClassifier, self).__init__()
        self.ignore_labels = rospy.get_param('~ignore_labels', [])
        self.target_names = rospy.get_param('~target_names', [])
        self.pub = self.advertise(
            '~output', ClassificationResult, queue_size=1)
 
    def subscribe(self):
        self.sub = rospy.Subscriber('~input', Image, self._cb)
 
    def unsubscribe(self):
        self.sub.unregister()
 
    def _cb(self, imgmsg):
        bridge = cv_bridge.CvBridge()
        img = bridge.imgmsg_to_cv2(imgmsg)
        label, proba = self._classify(img)
        msg = ClassificationResult()
        msg.header = imgmsg.header
        msg.labels = [label]
        msg.label_names = [self.target_names[label]]
        msg.label_proba = [proba[label]]
        msg.probabilities = proba
        msg.classifier = self.classifier_name
        msg.target_names = self.target_names
        self.pub.publish(msg)
 
    def _classify(self, label_img):
        counts = np.bincount(label_img.flatten(),
                             minlength=len(self.target_names))
        counts[self.ignore_labels] = 0
        label = np.argmax(counts)
        proba = counts.astype(np.float32) / counts.sum()
        return label, proba