split_fore_background.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 # -*- coding: utf-8 -*-
00003 
00004 import numpy as np
00005 
00006 import cv_bridge
00007 from jsk_recognition_utils.depth import split_fore_background
00008 from jsk_topic_tools import ConnectionBasedTransport
00009 import rospy
00010 from sensor_msgs.msg import Image
00011 
00012 
00013 class SplitForeBackground(ConnectionBasedTransport):
00014 
00015     def __init__(self):
00016         super(SplitForeBackground, self).__init__()
00017         self.fg_mask_pub_ = self.advertise(
00018             '~output/fg_mask', Image, queue_size=10)
00019         self.bg_mask_pub_ = self.advertise(
00020             '~output/bg_mask', Image, queue_size=10)
00021 
00022     def subscribe(self):
00023         self.sub = rospy.Subscriber('~input', Image, self._apply)
00024 
00025     def unsubscribe(self):
00026         self.sub.unregister()
00027 
00028     def _apply(self, depth_msg):
00029         # validation
00030         supported_encodings = {'16UC1', '32FC1'}
00031         if depth_msg.encoding not in supported_encodings:
00032             rospy.logwarn('Unsupported depth image encoding: {0}'
00033                           .format(depth_msg.encoding))
00034         # split fg/bg and get each mask
00035         bridge = cv_bridge.CvBridge()
00036         depth = bridge.imgmsg_to_cv2(depth_msg)
00037         nan_mask = None
00038         if depth_msg.encoding == '32FC1':
00039             # convert float to int (handle 32FC1 as 16UC1)
00040             nan_mask = np.isnan(depth)
00041             depth[nan_mask] = 0
00042             depth *= 255
00043             depth = depth.astype(np.uint8)
00044         fg_mask, bg_mask = split_fore_background(depth)
00045         if nan_mask is not None:
00046             fg_mask[nan_mask] = 0
00047             bg_mask[nan_mask] = 0
00048         # fg_mask
00049         fg_mask = (fg_mask * 255).astype(np.uint8)
00050         fg_mask_msg = bridge.cv2_to_imgmsg(fg_mask, encoding='mono8')
00051         fg_mask_msg.header = depth_msg.header
00052         self.fg_mask_pub_.publish(fg_mask_msg)
00053         # bg_mask
00054         bg_mask = (bg_mask * 255).astype(np.uint8)
00055         bg_mask_msg = bridge.cv2_to_imgmsg(bg_mask, encoding='mono8')
00056         bg_mask_msg.header = depth_msg.header
00057         self.bg_mask_pub_.publish(bg_mask_msg)
00058 
00059 
00060 if __name__ == '__main__':
00061     rospy.init_node('split_fore_background')
00062     split_fbg = SplitForeBackground()
00063     rospy.spin()


jsk_perception
Author(s): Manabu Saito, Ryohei Ueda
autogenerated on Sun Oct 8 2017 02:43:23