probability_image_classifier.py

Go to the documentation of this file.

#!/usr/bin/env python
 
from jsk_perception.label_image_classifier import LabelImageClassifier
import numpy as np
import rospy
 
 
class ProbabilityImageClassifier(LabelImageClassifier):
 
    classifier_name = 'probability_image_classifier'
 
    def __init__(self):
        super(ProbabilityImageClassifier, self).__init__()
 
    def _classify(self, proba_img):
        proba = np.nansum(proba_img, axis=(0, 1)).astype(np.float32)
 
        self.ignore_labels = np.asarray(self.ignore_labels)
 
        # validation
        assert proba.ndim == 1
        n_labels = proba.shape[0]
        mask_valid = self.ignore_labels < n_labels
        if mask_valid.sum() != len(self.ignore_labels):
            rospy.logwarn_throttle(
                10, "The max label value in '~ignore_labels' exceeds "
                    "the number of labels of input probability image.")
 
        # filtering unrelated label probabilities
        ignore_labels = self.ignore_labels[mask_valid]
        proba[ignore_labels] = 0
        label = np.argmax(proba)
        proba = proba / proba.sum()
 
        return label, proba
 
 
if __name__ == '__main__':
    rospy.init_node('probability_image_classifier')
    app = ProbabilityImageClassifier()
    rospy.spin()