shirt_color.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on 7/29/15

@author: sampfeiffer

shirt_color.py contains...
"""
__author__ = 'sampfeiffer'

import rospy
from sensor_msgs.msg import Image
from pal_detection_msgs.msg import FaceDetections, FaceDetection
from std_msgs.msg import String
from cv_bridge import CvBridge, CvBridgeError
import cv2
import numpy as np

IMAGE_TOPIC = "/xtion/rgb/image_rect_color"
FACES_TOPIC = "/pal_face/faces"
SHIRT_COLOR_TOPIC = "/shirt_color"
SHIRT_COLOR_DEBUG_TOPIC = "/shirt_color_debug"

# used http://colorizer.org/ and adapted the hsv colors with an ipython doing HUEVAL / 360 * 255
color_dict = {'red' : 0,
              'yellow' : 42,
              'green': 80,
              'blue': 155,
              'purple': 182,
              'pink': 218}
colors = ['red', 'yellow', 'green', 'blue', 'purple', 'pink']
values = [0, 42, 80, 155, 182, 218]
# saturation must be over 25%
# brightness must be over 50%


class ShirtColor(object):

    colors = ['red', 'yellow', 'green', 'blue', 'purple', 'pink']
    bgr_colors = [(0,0,255), (0,255,255), (0, 255, 0), (255, 0, 0), (153, 0, 153), (0, 255, 255)]
    colors_scale = ['red', 'yellow', 'green', 'blue', 'purple', 'pink', 'red']
    color_hsv_values = [0, 42, 80, 125, 182, 218, 255]

    def __init__(self):
        rospy.loginfo("Setting subscriber to: " + SHIRT_COLOR_TOPIC)
        self.shirt_color_pub = rospy.Publisher(SHIRT_COLOR_TOPIC, String, queue_size=1)
        self.bridge = CvBridge()
        rospy.loginfo("Setting publisher to: " + SHIRT_COLOR_DEBUG_TOPIC)
        self.shirt_color_debug_pub = rospy.Publisher(SHIRT_COLOR_DEBUG_TOPIC, Image, queue_size=1)
        rospy.loginfo("Subscribing to: " + IMAGE_TOPIC)
        self.last_img = None
        self.imgs_sub = rospy.Subscriber(IMAGE_TOPIC, Image, self.img_cb, queue_size=1)
        rospy.loginfo("Subscribing to: " + FACES_TOPIC)
        self.last_faces = None
        self.faces_sub = rospy.Subscriber(FACES_TOPIC, FaceDetections, self.faces_cb, queue_size=1)


    def img_cb(self, data):
        """
        :type data: Image
        """
        self.last_img = data


    def get_centered_face(self, faces, image_center_x, image_center_y):
        """
        :type faces:
        :return:
        """
        # Get the face that is more centered of the image if many faces
        centered_face = None
        for face_detection in faces:
            #face_detection = FaceDetection() # used for type completion in PyCharm
            if centered_face is None:
                centered_face = face_detection
            else:
                if (abs(image_center_x - face_detection.x) <= abs(image_center_x - centered_face.x)) \
                        and \
                        (abs(image_center_y - face_detection.y) <= abs(image_center_y - centered_face.y)):
                    centered_face = face_detection
        return centered_face

    def get_closer_color(self, hue_value):
        #print "getting closer color to: " + str(hue_value)
        closer_value = min(range(len(self.color_hsv_values)), key=lambda i: abs(self.color_hsv_values[i]-hue_value))
        return self.colors_scale[closer_value]


    # def get_average_9_pixels(self, image, x, y):
    #     avg_hue = ((image[x-1, y-1][0] + image[x, y-1][0] + image[x+1, y-1][0]) + \
    #               (image[x-1, y][0] + image[x, y][0] + image[x+1, y][0]) + \
    #               (image[x-1, y+1][0] + image[x, y+1][0] + image[x+1, y+1][0]))
    #     print avg_hue
    #     avg_hue = int(avg_hue / 9)
    #     avg_sat = ((image[x-1, y-1][1] + image[x, y-1][1] + image[x+1, y-1][1]) + \
    #               (image[x-1, y][1] + image[x, y][1] + image[x+1, y][1]) + \
    #               (image[x-1, y+1][1] + image[x, y+1][1] + image[x+1, y+1][1]))
    #     print avg_sat
    #     avg_sat = int(avg_sat / 9)
    #     avg_bri = ((image[x-1, y-1][2] + image[x, y-1][2] + image[x+1, y-1][2]) + \
    #               (image[x-1, y][2] + image[x, y][2] + image[x+1, y][2]) + \
    #               (image[x-1, y+1][2] + image[x, y+1][2] + image[x+1, y+1][2]))
    #     print avg_bri
    #     avg_bri = int(avg_bri / 9)
    #     return avg_hue, avg_sat, avg_bri


    def get_average_block(self, image_block):
        colors_found = []
        #print "image_block looks like: " + str(image_block)
        for row in image_block:
            #print "row is: " + str(row)
            for pixel in row:
                #print "pixel is: " + str(pixel)
                if pixel[1] > (0.25 * 255) and pixel[2] > (0.25 * 255):
                    this_color = self.get_closer_color(pixel[0])
                    colors_found.append(this_color)

        # get the color repeated most times:
        rep_times = []
        for col in self.colors:
            this_count = colors_found.count(col)
            rep_times.append(this_count)

        max_n = max(rep_times)
        indx = rep_times.index(max_n)
        if max_n > 0:
            return colors[indx]
        else:
            return None



    def get_bgr_from_color_name(self, color_name):
        return self.bgr_colors[self.colors.index(color_name)]

    def get_shirt_color(self, face_detection, image):
        """
        :type face_detection: FaceDetection
        :type image: Image
        """
        pixel_color_string = None
        # Get some pixel group bottom of the face
        # use height to get an idea of how close the person is to know how many pixels down we should look
        pixels_down = face_detection.height

        # get its color
        cvImage = self.bridge.imgmsg_to_cv2(image, "bgr8")
        hsv = cv2.cvtColor(cvImage, cv2.COLOR_BGR2HSV)


        block_x = face_detection.x + face_detection.width/2
        block_y = face_detection.y + face_detection.height/2 + pixels_down
        w = face_detection.width
        h = face_detection.height
        pixel_color_string = self.get_average_block(hsv[block_x-w/4:block_x+w/4, block_y-h/4:block_y+h/4])
        if pixel_color_string is None:
            return

        # publish debug image
        if self.shirt_color_debug_pub.get_num_connections() > 0:
            x = face_detection.x + face_detection.width/2
            y = face_detection.y + face_detection.height/2
            width = face_detection.width
            height = face_detection.height
            #center = (x, y + pixels_down)
            color = (0, 0, 250)

            # rectangle of the face detection
            p1 = (x - width/2, y - height/2)
            p2 = (x + width/2, y + height/2)
            cv2.rectangle(cvImage, p1, p2, color, thickness=3)

            interesting_color = self.get_bgr_from_color_name(pixel_color_string)
            # Rectangle of the area used for computing the color
            p1 = (block_x-w/4, block_y-h/4)
            p2 = (block_x+w/4, block_y+h/4)
            cv2.rectangle(cvImage, p1, p2, interesting_color, thickness=2)

            # text in the rectangle of the color computed
            text_orig = (x-w/4, y + pixels_down)
            cv2.putText(cvImage, pixel_color_string, text_orig, cv2.FONT_HERSHEY_SIMPLEX, 0.9, interesting_color, 1)
            self.shirt_color_debug_pub.publish(self.bridge.cv2_to_imgmsg(cvImage, "bgr8"))

        # return that color
        return pixel_color_string

    def faces_cb(self, data):
        """
        :type data: FaceDetections
        """
        # Don't do anything until we got an image
        if self.last_img is None:
            return

        centered_face = self.get_centered_face(data.faces, self.last_img.width/2, self.last_img.height/2)
        if centered_face is not None:
            shirt_color = self.get_shirt_color(centered_face, self.last_img)
            if shirt_color is not None:
                self.shirt_color_pub.publish(String(shirt_color))
                print "Found that the shirt is of color: " + str(shirt_color)



if __name__ == '__main__':
    rospy.init_node('shirt_color_node')
    sc = ShirtColor()
    rospy.spin()