diff --git a/people_recognition_2d/src/people_recognition_2d/people_recognizer_2d.py b/people_recognition_2d/src/people_recognition_2d/people_recognizer_2d.py index 190680d..ab02c05 100644 --- a/people_recognition_2d/src/people_recognition_2d/people_recognizer_2d.py +++ b/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)