facenet_pytorch final implementation

.devcontainer/devcontainer.json

@@ -0,0 +1,26 @@
+{
+  "name": "noetic desktop",
+  "image": "ghcr.io/tue-robotics/tue-env-ros-noetic:latest",
+  "runArgs": [
+    "--network=host",
+    "--gpus", "all"
+  ],
+  "workspaceMount": "source=${localWorkspaceFolder},target=/home/docker/ros/noetic/repos/github.com/tue-robotics/${localWorkspaceFolderBasename},type=bind",
+  "workspaceFolder": "/home/docker/ros/noetic/repos/github.com/tue-robotics/${localWorkspaceFolderBasename}",
+  "containerEnv": {
+        "DISPLAY": "unix:0"
+    },
+  "mounts": [
+    "source=/tmp/.X11-unix,target=/tmp/.X11-unix,type=bind,consistency=cached",
+    "source=${localEnv:HOME}/.bash_history,target=/home/docker/.bash_history,type=bind"
+  ],
+  "customizations" : {
+    "jetbrains" : {
+      "backend" : "IntelliJ"
+    }
+  },
+
+  "postCreateCommand": "tue-env config ros-noetic git-use-ssh; tue-get install git-ps1",
+
+  "postAttachCommand": "tue-get install ros-image_recognition_face_recognition",
+}

image_recognition_face_recognition/README.md

@@ -2,12 +2,12 @@
 
 TO DO
 
-Face recognition with use of Facenet (https://github.com/timesler/facenet-pytorch/tree/master)
+Face recognition with use of Facenet (<https://github.com/timesler/facenet-pytorch/tree/master>)
 Paper can be found: [paper](https://arxiv.org/pdf/1503.03832.pdf)
 
 ## Installation
 
-See https://github.com/tue-robotics/image_recognition
+See <https://github.com/tue-robotics/image_recognition>
 
 Make sure you have installed **CUDA8**.
 
@@ -24,19 +24,19 @@ Run the command why using your camera as an input (use roscamera node):
 
 ### ROS Node
 
-Run the image_recognition_face_detection node in one terminal (Specify the dlib and openface_net path as ROS parameter):
+Run the image_recognition_face_detection node in one terminal:
 
     rosrun image_recognition_face_detection face_recognition_node
 
-Run the rqt annotation client (https://github.com/tue-robotics/image_recognition_rqt)
+Run the rqt annotation client (<https://github.com/tue-robotics/image_recognition_rqt>)
 
     rosrun image_recognition_rqt annotation_gui
 
 Setup the service by clicking the gear wheel in the top-right corner. Select the `/annotate` services exposed by the openface ros node. Also don't forget to set-up your labels.
 
 ![Configuration](doc/config.png)
 
-Now draw a rectangle around the face you would like to learn. The face recognizer will find the biggest face in the image and store a representation for this face. 
+Now draw a rectangle around the face you would like to learn. The face recognizer will find the biggest face in the image and store a representation for this face.
 
 ![Annotate](doc/annotate.png)
 
@@ -46,7 +46,7 @@ Now select the label and you will see that the openface ros node stores the face
 
 Learn as many faces as you want ..
 
-Next step is starting the image_recognition_Rqt test gui (https://github.com/tue-robotics/image_recognition_rqt)
+Next step is starting the image_recognition_Rqt test gui (<https://github.com/tue-robotics/image_recognition_rqt>)
 
     rosrun image_recognition_rqt test_gui
 
@@ -56,7 +56,6 @@ Again configure the service you want to call with the gear-wheel in the top-righ
 
 You will see that the result of the detection will prompt in a dialog combo box. Also the detections will be drawn on the image.
 
-
 ### Command line
 
 Command line interface to test the detection / recognition based on an image:
@@ -67,4 +66,4 @@ Run the command on an example image:
 
     rosrun image_recognition_face_recognition get_face_recognition `rospack find image_recognition_face_recognition`/doc/example.png
 
-![Example](doc/example.png)
+![Example](doc/example.png)

image_recognition_face_recognition/package.xml

@@ -20,7 +20,7 @@
   <exec_depend>diagnostic_updater</exec_depend>
   <exec_depend>image_recognition_msgs</exec_depend>
   <exec_depend>image_recognition_util</exec_depend>
-  <exec_depend>python-dlib</exec_depend> <!--pip package-->
+  <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>

image_recognition_face_recognition/scripts/face_recognition_node

@@ -2,27 +2,22 @@
 
 import os
 import sys
+
 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, AnnotateRequest
 from image_recognition_util import image_writer
-from image_recognition_face_recognition.facenet_recognition import Facenet_recognition
+from sensor_msgs.msg import Image, RegionOfInterest
+from std_srvs.srv import Empty, EmptyRequest, EmptyResponse
+
+from image_recognition_face_recognition.face_recognizer import FaceRecognizer
 
 
 class OpenfaceROS:
     def __init__(
         self,
-        align_path,
-        net_path,
         save_images_folder,
         topic_save_images,
         service_save_images,
@@ -32,8 +27,6 @@ class OpenfaceROS:
         """
         OpenfaceROS class that wraps the FaceRecognizer in a ROS node
 
-        :param align_path: dlib face align path
-        :param net_path: openface neural net path
         :param save_images_folder: path where to store the images
         :param topic_save_images: whether to save images originated from image topic callback
         :param service_save_images: whether to save images originated from a service call
@@ -42,7 +35,7 @@ class OpenfaceROS:
         """
 
         # Openface ROS
-        self._face_recognizer = Facenet_recognition()
+        self._face_recognizer = FaceRecognizer()
         self._save_images_folder = save_images_folder
         self._topic_save_images = topic_save_images
         self._service_save_images = service_save_images
@@ -65,15 +58,13 @@ class OpenfaceROS:
             self._result_image_publisher = rospy.Publisher("result_image", Image, queue_size=10)
 
         rospy.loginfo("OpenfaceROS initialized:")
-        rospy.loginfo(" - dlib_align_path=%s", align_path)
-        rospy.loginfo(" - openface_net_path=%s", net_path)
         rospy.loginfo(" - save_images_folder=%s", save_images_folder)
         rospy.loginfo(" - topic_save_images=%s", topic_save_images)
         rospy.loginfo(" - service_save_images=%s", service_save_images)
         rospy.loginfo(" - topic_publish_result_image=%s", topic_publish_result_image)
         rospy.loginfo(" - service_publish_result_image=%s", service_publish_result_image)
 
-    def _annotate_srv(self, req):
+    def _annotate_srv(self, req: AnnotateRequest):
         # Convert to opencv image
         """
         Annotate service callback that trains the face of the user
@@ -84,12 +75,12 @@ class OpenfaceROS:
         try:
             bgr_image = self._bridge.imgmsg_to_cv2(req.image, "bgr8")
         except CvBridgeError as e:
-            raise Exception("Could not convert to opencv image: %s" % str(e))
+            raise RuntimeError(f"Could not convert to opencv image: {e}")
 
         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:
@@ -98,87 +89,76 @@ class OpenfaceROS:
             try:
                 self._face_recognizer.train(roi_image, annotation.label)
             except Exception as e:
-                raise Exception("Could not get representation of face image: %s" % str(e))
+                raise RuntimeError(f"Could not get representation of face image: {e}")
 
             rospy.loginfo("Succesfully learned face of '%s'" % annotation.label)
 
-        return {}
-
-    def _clear_srv(self, req):
+    def _clear_srv(self, _: EmptyRequest):
         """
         Service to clear the trained faces
-
-        :param req: Empty
-        :return: Empty
         """
         self._face_recognizer.clear_trained_faces()
-        return {}
+        return EmptyResponse()
 
-    def _get_recognitions(self, image_msg, save_images, publish_images):
+    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 RuntimeError(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 = []
-        # Added this line for testing with some labels
-        labels = ["jason", "kona"]
-
-        # Call facebet neural network in two stages
-        face_recognitions = self._face_recognizer.face_detection(bgr_image)
-        distance, labels_pred = self._face_recognizer.detection_recognition(bgr_image, labels, train=True)
-
+
+        # 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:
-            import math
-
-            face_recognition = [math.floor(xi) for xi in fr]
+            label = fr.l2_distances[0].label if fr.l2_distances else "face_unknown"
+
             if save_images:
-                label = labels_pred[label_idx]
+                roi = fr.roi
                 roi_image = bgr_image[
-                    face_recognition[2]: face_recognition[3],
-                    face_recognition[0]: face_recognition[1],
-                ]
-                image_writer.write_annotated(self._save_images_folder, roi_image, label, False)
+                            roi.y_offset:roi.y_offset + roi.height,
+                            roi.x_offset:roi.x_offset + roi.width]
+                image_writer.write_annotated(self._save_images_folder, fr.image, label, False)
                 images.append(roi_image)
                 labels.append(label)
-            label = labels_pred[label_idx]
-            distance_fr = distance[label_idx]
+
             recognitions.append(
                 Recognition(
                     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=l2.label, probability=1.0 / (l2.distance + 0.001))
+                            for l2 in fr.l2_distances
                         ],
                     ),
                     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,
                     ),
                 )
             )
-            label_idx = label_idx + 1
 
         annotated_original_image = image_writer.get_annotated_cv_image(bgr_image, recognitions)
         if save_images:
@@ -197,7 +177,7 @@ class OpenfaceROS:
         return recognitions
 
     def _image_callback(self, image_msg):
-        # Comment this exception for beeter debbuging
+        # Comment this exception for better debuging
         try:
             recognitions = self._get_recognitions(
                 image_msg,
@@ -220,22 +200,10 @@ class OpenfaceROS:
         # Service response
         return {"recognitions": recognitions}
 
-    def restore_trained_faces(self, file_name):
-        self._face_recognizer.restore_trained_faces(file_name)
-
 
 if __name__ == "__main__":
     rospy.init_node("face_recognition")
     try:
-        dlib_shape_predictor_path = os.path.expanduser(
-            rospy.get_param(
-                "~align_path",
-                "~/openface/models/dlib/shape_predictor_68_face_landmarks.dat",
-            )
-        )
-        openface_neural_network_path = os.path.expanduser(
-            rospy.get_param("~net_path", "~/openface/models/openface/nn4.small2.v1.t7")
-        )
         save_images = rospy.get_param("~save_images", True)
         topic_save_images = rospy.get_param("~topic_save_images", False)
         service_save_images = rospy.get_param("~service_save_images", True)
@@ -248,22 +216,23 @@ if __name__ == "__main__":
 
         save_images_folder = None
         if save_images:
-            save_images_folder = os.path.expanduser(rospy.get_param("~save_images_folder", "/tmp/openface"))
+            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)
 
     image_recognition_openface = OpenfaceROS(
-        dlib_shape_predictor_path,
-        openface_neural_network_path,
         save_images_folder,
         topic_save_images,
         service_save_images,
         topic_publish_result_image,
         service_publish_result_image,
     )
+
     if db:
-        rospy.loginfo("loading face database from %s", db)
+        rospy.loginfo(f"loading face database from {db}")
         image_recognition_openface.restore_trained_faces(db)
 
     updater = diagnostic_updater.Updater()