sound_classifier.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 
3 # Copied from https://github.com/jsk-ros-pkg/jsk_recognition/blob/master/jsk_perception/node_scripts/vgg16_object_recognition.py
4 
5 from __future__ import absolute_import
6 from __future__ import division
7 from __future__ import print_function
8 
9 import chainer
10 from chainer import cuda
11 import chainer.serializers as S
12 from chainer import Variable
13 from distutils.version import LooseVersion
14 import numpy as np
15 import skimage.transform
16 
17 import cv_bridge
18 from jsk_recognition_msgs.msg import ClassificationResult
19 from sound_classification.nin.nin import NIN
20 from sound_classification.vgg16.vgg16_batch_normalization import VGG16BatchNormalization
21 from jsk_topic_tools import ConnectionBasedTransport # TODO use LazyTransport
22 import os.path as osp
24 import rospy
25 from sensor_msgs.msg import Image
26 
27 from train import PreprocessedDataset
28 
29 
30 class SoundClassifier(ConnectionBasedTransport):
31  """
32  Classify spectrogram using neural network
33  input: sensor_msgs/Image, 8UC1
34  output jsk_recognition_msgs/ClassificationResult
35  """
36 
37  def __init__(self):
38  super(self.__class__, self).__init__()
39  self.gpu = rospy.get_param('~gpu', -1)
41  self.target_names_ordered = self.dataset.target_classes
42  self.target_names = rospy.get_param('~target_names', self.target_names_ordered)
43  for i, name in enumerate(self.target_names):
44  if not name.endswith('\n'):
45  self.target_names[i] = name + '\n'
46  self.model_name = rospy.get_param('~model_name')
47  if self.model_name == 'nin':
48  self.insize = 227
49  self.model = NIN(n_class=len(self.target_names))
50  elif self.model_name == 'vgg16':
51  self.insize = 224
53  n_class=len(self.target))
54  else:
55  rospy.logerr('Unsupported ~model_name: {0}'
56  .format(self.model_name))
57  model_file = rospy.get_param(
58  '~model_file',
59  osp.join(self.dataset.root, 'result',
60  self.model_name, 'model_best.npz'))
61  S.load_npz(model_file, self.model)
62  if self.gpu != -1:
63  self.model.to_gpu(self.gpu)
64  self.pub = self.advertise('~output', ClassificationResult,
65  queue_size=1)
66  self.pub_input = self.advertise(
67  '~debug/net_input', Image, queue_size=1)
68 
69  def subscribe(self):
70  sub = rospy.Subscriber(
71  '~input', Image, self._recognize, callback_args=None,
72  queue_size=1, buff_size=2**24)
73  self.subs = [sub]
74 
75  def unsubscribe(self):
76  for sub in self.subs:
77  sub.unregister()
78 
79  def _recognize(self, imgmsg):
80  bridge = cv_bridge.CvBridge()
81  mono = bridge.imgmsg_to_cv2(imgmsg)
82  bgr = img_jet(mono)
83  bgr = skimage.transform.resize(
84  bgr, (self.insize, self.insize), preserve_range=True)
85  input_msg = bridge.cv2_to_imgmsg(bgr.astype(np.uint8), encoding='bgr8')
86  input_msg.header = imgmsg.header
87  self.pub_input.publish(input_msg)
88 
89  # (Height, Width, Channel) -> (Channel, Height, Width)
90  # ###
91  # print(type())
92  # print(rgb.shape)
93  # import cv2
94  # cv2.imwrite('/home/naoya/test.png', rgb)
95  # ###
96  rgb = bgr.transpose((2, 0, 1))[::-1, :, :]
97  rgb = self.dataset.process_image(rgb)
98  x_data = np.array([rgb], dtype=np.float32)
99  if self.gpu != -1:
100  x_data = cuda.to_gpu(x_data, device=self.gpu)
101  if LooseVersion(chainer.__version__) < LooseVersion('2.0.0'):
102  x = Variable(x_data, volatile=True)
103  self.model.train = False
104  self.model(x)
105  else:
106  with chainer.using_config('train', False), \
107  chainer.no_backprop_mode():
108  x = Variable(x_data)
109  self.model(x)
110 
111  # swap_labels[label number in self.target_names]
112  # -> label number in self.target_names_ordered
113  swap_labels = [self.target_names_ordered.index(name) for name in self.target_names]
114  for i in range(len(swap_labels)):
115  if not (i in swap_labels):
116  rospy.logerr('Wrong target_names is given by rosparam.')
117  exit()
118  proba = cuda.to_cpu(self.model.pred.data)[0]
119  proba_swapped = [proba[swap_labels[i]] for i, p in enumerate(proba)]
120  label_swapped = np.argmax(proba_swapped)
121  label_name = self.target_names[label_swapped]
122  label_proba = proba_swapped[label_swapped]
123  cls_msg = ClassificationResult(
124  header=imgmsg.header,
125  labels=[label_swapped],
126  label_names=[label_name],
127  label_proba=[label_proba],
128  probabilities=proba_swapped,
129  classifier=self.model_name,
130  target_names=self.target_names,
131  )
132  self.pub.publish(cls_msg)
133 
134 
135 if __name__ == '__main__':
136  rospy.init_node('sound_classifier')
138  rospy.spin()
sound_classification.vgg16.vgg16_batch_normalization
Definition: vgg16_batch_normalization.py:1
sound_classifier.SoundClassifier.model
model
Definition: sound_classifier.py:49
sound_classifier.SoundClassifier.pub_input
pub_input
Definition: sound_classifier.py:66
sound_classifier.SoundClassifier.insize
insize
Definition: sound_classifier.py:48
sound_classification.nin.nin.NIN
Definition: nin.py:9
sound_classifier.SoundClassifier.subscribe
def subscribe(self)
Definition: sound_classifier.py:69
sound_classifier.SoundClassifier._recognize
def _recognize(self, imgmsg)
Definition: sound_classifier.py:79
sound_classification.process_gray_image
Definition: process_gray_image.py:1
sound_classification.vgg16.vgg16_batch_normalization.VGG16BatchNormalization
Definition: vgg16_batch_normalization.py:8
sound_classifier.SoundClassifier.dataset
dataset
Definition: sound_classifier.py:40
sound_classifier.SoundClassifier.pub
pub
Definition: sound_classifier.py:64
sound_classification.process_gray_image.img_jet
def img_jet(img)
Definition: process_gray_image.py:39
sound_classifier.SoundClassifier.target_names
target_names
Definition: sound_classifier.py:42
sound_classifier.SoundClassifier.model_name
model_name
Definition: sound_classifier.py:46
sound_classifier.SoundClassifier.subs
subs
Definition: sound_classifier.py:73
sound_classifier.SoundClassifier.target_names_ordered
target_names_ordered
Definition: sound_classifier.py:41
sound_classifier.SoundClassifier
Definition: sound_classifier.py:30
sound_classification.nin.nin
Definition: nin.py:1
sound_classifier.SoundClassifier.__init__
def __init__(self)
Definition: sound_classifier.py:37
sound_classifier.SoundClassifier.unsubscribe
def unsubscribe(self)
Definition: sound_classifier.py:75
sound_classifier.SoundClassifier.gpu
gpu
Definition: sound_classifier.py:39
train.PreprocessedDataset
Definition: train.py:29


sound_classification
Author(s):
autogenerated on Fri May 16 2025 03:12:55