Add hair roi to people_recognition

people_recognition_2d/src/people_recognition_2d/people_recognizer_2d.py

@@ -15,6 +15,8 @@
 from rospy import ServiceException
 from sensor_msgs.msg import Image, RegionOfInterest
 
+import cv2
+
 
 def _get_and_wait_for_service(srv_name, srv_class):
     """
@@ -219,6 +221,19 @@ def _shirt_colors_to_label(shirt_colors):
             label += " {}".format(color)
         return label
 
+    @staticmethod
+    def _hair_colors_to_label(hair_colors):
+        """
+        Convert shirt colors array to label string
+
+        :param: shirt_colors: Array to colors
+        :return: string label
+        """
+        label = " hair colors:"
+        for color in hair_colors:
+            label += " {}".format(color)
+        return label
+
     @staticmethod
     def _shirt_roi_from_face_roi(face_roi, image_shape):
         """
@@ -236,6 +251,22 @@ def _shirt_roi_from_face_roi(face_roi, image_shape):
         rospy.logdebug("face_roi: {}, shirt_roi: {}, img.shape: {}".format(face_roi, shirt_roi, image_shape))
         return shirt_roi
 
+    @staticmethod
+    def _hair_roi_from_face_roi(face_roi, image_shape):
+        """
+        Given a ROI for a face, shift the ROI to the person's hair. Assuming the person is upright :/
+
+        :param face_roi: RegionOfInterest
+        :param image_shape: tuple of the image shape
+        :return: RegionOfInterest
+        """
+        hair_roi = copy.deepcopy(face_roi)
+        hair_roi.height = max(int(face_roi.height/6), 5)
+        hair_roi.width = max(face_roi.width, 5)
+        hair_roi.y_offset = min(hair_roi.y_offset, image_shape[0] - hair_roi.height)
+        rospy.logdebug("face_roi: {}, hair_roi: {}, img.shape: {}".format(face_roi, hair_roi, image_shape))
+        return hair_roi
+
     def recognize(self, image_msg):
         assert isinstance(image_msg, Image)
         image = self._bridge.imgmsg_to_cv2(image_msg)
@@ -282,30 +313,42 @@ def recognize(self, image_msg):
         if self._enable_shirt_color_extraction:
             rospy.loginfo("_get_color_extractor...")
             shirt_colors_array = []
+            hair_colors_array = []
             for r in face_recognitions:
                 shirt_roi = PeopleRecognizer2D._shirt_roi_from_face_roi(r.roi, image.shape)
+                hair_roi = PeopleRecognizer2D._hair_roi_from_face_roi(r.roi, image.shape)
+
                 shirt_image_msg = self._bridge.cv2_to_imgmsg(PeopleRecognizer2D._image_from_roi(image, shirt_roi))
+                hair_image_msg = self._bridge.cv2_to_imgmsg(PeopleRecognizer2D._image_from_roi(image, hair_roi))
                 try:
-                    color_extractor_response = _get_service_response(self._color_extractor_srv, shirt_image_msg)
+                    color_shirt_extractor_response = _get_service_response(self._color_extractor_srv, shirt_image_msg)
+                    color_hair_extractor_response = _get_service_response(self._color_extractor_srv, hair_image_msg)
                 except ServiceException as e:
                     rospy.logerr("Color extractor service request failed: {}".format(e))
                     shirt_colors = []
+                    hair_colors = []
                 else:
                     shirt_colors = [p.label for p in
-                                    color_extractor_response.recognitions[0].categorical_distribution.probabilities]
+                                    color_shirt_extractor_response.recognitions[0].categorical_distribution.probabilities]
+                    hair_colors = [p.label for p in
+                                    color_hair_extractor_response.recognitions[0].categorical_distribution.probabilities]
                 shirt_colors_array.append(shirt_colors)
+                hair_colors_array.append(hair_colors)
         else:
             shirt_colors_array = [[]] * len(face_recognitions)
+            hair_colors_array = [[]] * len(face_recognitions)
 
         # Prepare image annotation labels and People message
-        for face_label, face_properties, shirt_colors, body_parts, face_recognition in zip(face_labels,
-                                                                                           face_properties_array,
-                                                                                           shirt_colors_array,
-                                                                                           body_parts_array,
-                                                                                           face_recognitions):
+        for face_label, face_properties, shirt_colors, hair_colors, body_parts, face_recognition in zip(face_labels,
+                                                                                                        face_properties_array,
+                                                                                                        shirt_colors_array,
+                                                                                                        hair_colors_array,
+                                                                                                        body_parts_array,
+                                                                                                        face_recognitions):
 
             temp_label = PeopleRecognizer2D._face_properties_to_label(face_properties) + \
-                         PeopleRecognizer2D._shirt_colors_to_label(shirt_colors)
+                         PeopleRecognizer2D._shirt_colors_to_label(shirt_colors) + \
+                         PeopleRecognizer2D._hair_colors_to_label(hair_colors)
 
             if face_label:
                 image_annotations.append(face_label + " " + temp_label)