a crude color segmentation

image_converter.py

#!/usr/bin/env python

import sys
import rospy
import cv2
import numpy as np
# from std_msgs.msg import String
from sensor_msgs.msg import Image, PointCloud2
from cv_bridge import CvBridge, CvBridgeError

from detect_color import boundaries, test, everything


class VideoGetter:
    # obtain video message from the ROS topic specified and store the message content

    def __init__(self, video_topic, msg):
        """
        :param video_topic: string
        :param msg: object
        """
        self.output = None
        self.msg = msg
        self.video_topic = video_topic
        self.video_sub = rospy.Subscriber(video_topic, msg, self.video_callback)

    def video_callback(self, data):
        self.output = data


class ImageConverter(VideoGetter):
    # Convert ROS image messages to OpenCV images
    # the converted image is stored in self.output

    def __init__(self, input_topic, input_msg, output_topic=None, output_msg=None):
        VideoGetter.__init__(self, input_topic, input_msg)
        self.bridge = CvBridge()
        if (output_msg is not None) and (output_topic is not None):
            self.output_pub = rospy.Publisher(output_topic, output_msg, queue_size=0)

    def video_callback(self, data):
        # override with video conversion function
        # self.video_callback(self, data)
        try:
            self.output = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print e


class ColorsegProcessor(ImageConverter):
    # segment the colors based on input

    def __init__(self, input_topic, input_msg, output_topic=None, output_msg=None):
        ImageConverter.__init__(self, input_topic, input_msg, output_topic, output_msg)
        self.boundaries = self.obtain_boundaries()

    def obtain_boundaries(self):
        # TODO: Come up with better ways to extract color info and generate boundaries real time
        return self.output

    def video_callback(self, data):
        # convert the image to OpenCV format and store at self.output
        ImageConverter.video_callback(self, data)
        # calculate the mask
        # boundaries & test hard coded in a separate file
        overall_mask = cv2.inRange(self.output, test[0], test[1])
        for (l, u) in boundaries:
            overall_mask = np.add(cv2.inRange(self.output, l, u), overall_mask)
            # print overall_mask
        cv_image = cv2.bitwise_and(self.output, self.output, mask=overall_mask)
        try:
            self.output_pub.publish(self.bridge.cv2_to_imgmsg(cv_image, "bgr8"))
        except CvBridgeError as e:
            print e


def main(args):
    rospy.init_node('Color_segmentation', anonymous=True)
    segmenter = ColorsegProcessor('/camera/rgb/image_color', Image, '/processed_image', Image)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down" 
    # cv2.destroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)