[face] WIP

image_recognition_face_recognition/package.xml

@@ -20,9 +20,11 @@
   <exec_depend>diagnostic_updater</exec_depend>
   <exec_depend>image_recognition_msgs</exec_depend>
   <exec_depend>image_recognition_util</exec_depend>
+  <exec_depend>python3-facenet-pytorch-pip</exec_depend>
   <exec_depend>python3-numpy</exec_depend>
   <exec_depend>python3-opencv</exec_depend>
   <exec_depend>rospy</exec_depend>
+  <exec_depend>std_srvs</exec_depend>
 
   <test_depend>catkin_lint_cmake</test_depend>
 

image_recognition_face_recognition/scripts/face_recognition_node

@@ -2,19 +2,17 @@
 import math
 import os
 import sys
+from typing import TYPE_CHECKING
+
 import diagnostic_updater
 import rospy
 from cv_bridge import CvBridge, CvBridgeError
-from sensor_msgs.msg import RegionOfInterest, Image
-from std_srvs.srv import Empty
-from image_recognition_msgs.msg import (
-    Recognition,
-    Recognitions,
-    CategoryProbability,
-    CategoricalDistribution,
-)
-from image_recognition_msgs.srv import Recognize, Annotate
+from image_recognition_msgs.msg import CategoricalDistribution, CategoryProbability, Recognition, Recognitions
+from image_recognition_msgs.srv import Annotate, Recognize
 from image_recognition_util import image_writer
+from sensor_msgs.msg import Image, RegionOfInterest
+from std_srvs.srv import Empty, EmptyRequest, EmptyResponse
+
 from image_recognition_face_recognition.face_recognizer import FaceRecognizer
 
 
@@ -48,6 +46,8 @@ class OpenfaceROS:
         self._bridge = CvBridge()
         self._annotate_srv = rospy.Service("annotate", Annotate, self._annotate_srv)
         self._recognize_srv = rospy.Service("recognize", Recognize, self._recognize_srv)
+        self._clear_srv = rospy.Service("clear", Empty, self._clear_srv)
+
         self._image_subscriber = rospy.Subscriber("image", Image, self._image_callback)
         self._recognitions_publisher = rospy.Publisher("recognitions", Recognitions, queue_size=10)
 
@@ -80,8 +80,8 @@ class OpenfaceROS:
 
         for annotation in req.annotations:
             roi_image = bgr_image[
-                annotation.roi.y_offset: annotation.roi.y_offset + annotation.roi.height,
-                annotation.roi.x_offset: annotation.roi.x_offset + annotation.roi.width,
+                annotation.roi.y_offset : annotation.roi.y_offset + annotation.roi.height,
+                annotation.roi.x_offset : annotation.roi.x_offset + annotation.roi.width,
             ]
 
             if self._save_images_folder:
@@ -96,45 +96,53 @@ class OpenfaceROS:
 
         return {}
 
-    def _get_recognitions(self, image_msg, save_images, publish_images):
+    def _clear_srv(self, _: EmptyRequest):
+        """
+        Service to clear the trained faces
+        """
+        self._face_recognizer.clear_trained_faces()
+        return EmptyResponse()
+
+    def _get_recognitions(self, image_msg: Image, save_images: bool, publish_images: bool):
         # Convert to opencv image
         """
         Recognize service callback
 
-        :param req: The input image
+        :param image_msg: The input image
         :return: Recognitions
         """
         try:
             bgr_image = self._bridge.imgmsg_to_cv2(image_msg, "bgr8")
         except CvBridgeError as e:
-            raise Exception("Could not convert to opencv image: %s" % str(e))
+            raise Exception(f"Could not convert to opencv image: {e}")
 
         # Write raw image
         if save_images:
             image_writer.write_raw(self._save_images_folder, bgr_image)
 
         images = []
         labels = []
-        
-        # Call facebet neural network in two stages
-        face_recognitions = self._face_recognizer.face_detection(bgr_image)
+
+        # Call facenet neural network in two stages
+        face_recognitions = self._face_recognizer.detect(bgr_image)
         distance, labels_pred = self._face_recognizer.detection_recognition(bgr_image, labels, train=True)
-        
+
         # Fill recognitions
         recognitions = []
 
         # rospy.loginfo("Face recognitions: %s", face_recognitions)
 
         label_idx = 0
-        for fr in face_recognitions:            
+        for fr in face_recognitions:
             face_recognition = [math.floor(xi) for xi in fr]
             if save_images:
                 label = labels_pred[label_idx]
                 roi_image = bgr_image[
-                    face_recognition[2]: face_recognition[3],
-                    face_recognition[0]: face_recognition[1],
+                    face_recognition[2] : face_recognition[3],
+                    face_recognition[0] : face_recognition[1],
                 ]
                 image_writer.write_annotated(self._save_images_folder, roi_image, label, False)
+
                 images.append(roi_image)
                 labels.append(label)
             label = labels_pred[label_idx]
@@ -144,15 +152,16 @@ class OpenfaceROS:
                     categorical_distribution=CategoricalDistribution(
                         unknown_probability=0.0,  # TODO: When is it unknown?
                         probabilities=[
-                            # This line needs some changing 
-                            CategoryProbability(label=label, probability=1.0 / (distance_fr + 0.001)) for l2 in face_recognition
+                            # This line needs some changing
+                            CategoryProbability(label=label, probability=1.0 / (distance_fr + 0.001))
+                            for l2 in face_recognition
                         ],
                     ),
                     roi=RegionOfInterest(
-                        x_offset=face_recognition[0],
-                        y_offset=face_recognition[1],
-                        width=face_recognition[2] - face_recognition[0],
-                        height=face_recognition[3] - face_recognition[1],
+                        x_offset=face_recognition.roi.x_offset,
+                        y_offset=face_recognition.roi.y_offset,
+                        width=face_recognition.roi.width,
+                        height=face_recognition.roi.height,
                     ),
                 )
             )
@@ -175,7 +184,7 @@ class OpenfaceROS:
         return recognitions
 
     def _image_callback(self, image_msg):
-        # Comment this exception for beeter debbuging
+        # Comment this exception for better debbuging
         try:
             recognitions = self._get_recognitions(
                 image_msg,
@@ -208,9 +217,15 @@ if __name__ == "__main__":
         topic_publish_result_image = rospy.get_param("~topic_publish_result_image", True)
         service_publish_result_image = rospy.get_param("~service_publish_result_image", True)
 
+        db = rospy.get_param("~db", None)
+        if db:
+            db = os.path.expanduser(db)
+
         save_images_folder = None
         if save_images:
-            save_images_folder = os.path.expanduser(rospy.get_param("~save_images_folder", "/tmp/facenet_saved_images"))
+            save_images_folder = os.path.expanduser(
+                rospy.get_param("~save_images_folder", "/tmp/face_recognition_saved_images")
+            )
     except KeyError as e:
         rospy.logerr("Parameter %s not found" % e)
         sys.exit(1)
@@ -222,6 +237,11 @@ if __name__ == "__main__":
         topic_publish_result_image,
         service_publish_result_image,
     )
+
+    if db:
+        rospy.loginfo(f"loading face database from {db}")
+        image_recognition_openface.restore_trained_faces(db)
+
     updater = diagnostic_updater.Updater()
     updater.setHardwareID("none")
     updater.add(diagnostic_updater.Heartbeat())

image_recognition_face_recognition/scripts/get_face_recognition

@@ -1,14 +1,15 @@
 #!/usr/bin/env python
 
 import math
-from typing import Optional
+from typing import List, Optional
 
 import cv2
-from image_recognition_face_recognition.face_recognizer import FaceRecognizer
 from image_recognition_msgs.msg import Recognition
 from image_recognition_util import image_writer
 from sensor_msgs.msg import RegionOfInterest
 
+from image_recognition_face_recognition.face_recognizer import FaceRecognizer, RecognizedFace
+
 
 def main(image, db: Optional[str] = None):
     # Read the image
@@ -20,26 +21,21 @@ def main(image, db: Optional[str] = None):
     if db:
         face_recognizer.restore_trained_faces(db)
 
-    recognized_faces = face_recognizer.face_detection(img)
-    print(recognized_faces)
+    recognized_faces: List[RecognizedFace] = face_recognizer.detect(img)
 
     recognitions = []
     for fr in recognized_faces:
-        face_recognition = [math.floor(xi) for xi in fr]
         recognitions.append(
             Recognition(
                 roi=RegionOfInterest(
-                    x_offset=face_recognition[0],
-                    y_offset=face_recognition[1],
-                    width=face_recognition[2] - face_recognition[0],
-                    height=face_recognition[3] - face_recognition[1],
+                    x_offset=fr.roi.x_offset, y_offset=fr.roi.y_offset, width=fr.roi.width, height=fr.roi.height
                 )
             )
         )
 
     annotated_original_image = image_writer.get_annotated_cv_image(img, recognitions)
     cv2.imshow("result", annotated_original_image)
-    cv2.waitKey(1000)
+    cv2.waitKey()
 
 
 if __name__ == "__main__":

image_recognition_face_recognition/scripts/train_from_images

@@ -1,10 +1,11 @@
 #!/usr/bin/env python
+import logging
 import os
+from argparse import Action, ArgumentParser, ArgumentTypeError, FileType, Namespace
+
 import cv2
-import logging
-from argparse import ArgumentParser, Action, ArgumentTypeError, FileType
-from image_recognition_face_recognition.face_recognizer import FaceRecognizer
 
+from image_recognition_face_recognition.face_recognizer import FaceRecognizer
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
@@ -20,30 +21,30 @@ def main(modeldir, outfile, verbose: bool = False):
 
     dirs = [o for o in os.listdir(modeldir) if os.path.isdir(os.path.join(modeldir, o))]
     for cat in dirs:
-        logger.info('loading images for %s', cat)
+        logger.info("loading images for %s", cat)
         path = os.path.join(modeldir, cat)
 
-        logger.debug('loading images from %s', path)
+        logger.debug("loading images from %s", path)
         for filename in os.listdir(path):
             f = os.path.join(path, filename)
-            logger.debug('processing %s', f)
+            logger.debug("processing %s", f)
 
             img = cv2.imread(f, 1)  # load as color
 
             try:
-                logger.debug('training...')
+                logger.debug("training...")
                 face_recognizer.train(img, cat)
-                logger.debug('success!')
+                logger.debug("success!")
             except Exception:
-                logger.exception('face_recognizer failed to process %s', f)
+                logger.exception("face_recognizer failed to process %s", f)
 
-    logger.info('saving database...')
+    logger.info("saving database...")
     face_recognizer.save_trained_faces(outfile)
-    logger.info('done!')
+    logger.info("done!")
 
 
 class ReadableDir(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(self, parser: ArgumentParser, namespace: Namespace, values, option_string=None):
         prospective_dir = values
         if not os.path.isdir(prospective_dir):
             raise ArgumentTypeError("readable_dir:{0} is not a valid path".format(prospective_dir))
@@ -53,13 +54,13 @@ class ReadableDir(Action):
             raise ArgumentTypeError("readable_dir:{0} is not a readable dir".format(prospective_dir))
 
 
-if __name__ == '__main__':
-    parser = ArgumentParser(description='Train openface from a database of images')
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Train openface from a database of images")
 
-    parser.add_argument('modeldir', action=ReadableDir, help='Directory with folders for each category')
-    parser.add_argument('outfile', type=FileType('wb'), help='Where to output the trained faces database')
+    parser.add_argument("modeldir", action=ReadableDir, help="Directory with folders for each category")
+    parser.add_argument("outfile", type=FileType("wb"), help="Where to output the trained faces database")
 
-    parser.add_argument('-v', '--verbose', action='store_true')
+    parser.add_argument("-v", "--verbose", action="store_true")
 
     args = parser.parse_args()
     main(**vars(args))