From 8cb58c3249892d32137912de7bd75168f717c693 Mon Sep 17 00:00:00 2001
From: m-barker <mattbarker322@gmail.com>
Date: Mon, 5 Feb 2024 10:37:03 +0000
Subject: [PATCH] feat: publish collage similar faces

---
 .../vision/lasr_vision_deepface/nodes/service | 17 ++++++++++++--
 .../src/lasr_vision_deepface/deepface.py      | 22 +++++++++++--------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/common/vision/lasr_vision_deepface/nodes/service b/common/vision/lasr_vision_deepface/nodes/service
index 3b798b6cf..816d61249 100644
--- a/common/vision/lasr_vision_deepface/nodes/service
+++ b/common/vision/lasr_vision_deepface/nodes/service
@@ -30,15 +30,28 @@ if DEBUG:
 
 def detect(request: RecogniseRequest) -> RecogniseResponse:
     debug_publisher = None
+    similar_face_debug_publisher = None
     if DEBUG:
         if request.dataset in recognise_debug_publishers:
-            debug_publisher = recognise_debug_publishers[request.dataset]
+            debug_publisher = recognise_debug_publishers[request.dataset][0]
+            similar_face_debug_publisher = recognise_debug_publishers[request.dataset][
+                1
+            ]
         else:
             topic_name = re.sub(r"[\W_]+", "", request.dataset)
             debug_publisher = rospy.Publisher(
                 f"/recognise/debug/{topic_name}", Image, queue_size=1
             )
-    return face_recognition.detect(request, debug_publisher)
+            similar_face_debug_publisher = rospy.Publisher(
+                f"/recognise/debug/{topic_name}/similar_face", Image, queue_size=1
+            )
+            recognise_debug_publishers[request.dataset] = (
+                debug_publisher,
+                similar_face_debug_publisher,
+            )
+    return face_recognition.detect(
+        request, debug_publisher, similar_face_debug_publisher
+    )
 
 
 def learn_face(request: LearnFaceRequest) -> LearnFaceResponse:
diff --git a/common/vision/lasr_vision_deepface/src/lasr_vision_deepface/deepface.py b/common/vision/lasr_vision_deepface/src/lasr_vision_deepface/deepface.py
index af233a848..564a1cab1 100644
--- a/common/vision/lasr_vision_deepface/src/lasr_vision_deepface/deepface.py
+++ b/common/vision/lasr_vision_deepface/src/lasr_vision_deepface/deepface.py
@@ -48,6 +48,7 @@ def create_image_collage(images, output_size=(640, 480)):
     # Calculate grid dimensions
     num_images = len(images)
     rows = int(np.sqrt(num_images))
+    print(num_images, rows)
     cols = (num_images + rows - 1) // rows  # Ceiling division
 
     # Resize images to fit in the grid
@@ -112,7 +113,9 @@ def create_dataset(
 
 
 def detect(
-    request: RecogniseRequest, debug_publisher: rospy.Publisher | None
+    request: RecogniseRequest,
+    debug_publisher: rospy.Publisher | None,
+    debug_inference_pub: rospy.Publisher | None,
 ) -> RecogniseResponse:
     # Decode the image
     rospy.loginfo("Decoding")
@@ -134,9 +137,6 @@ def detect(
     except ValueError:
         return response
 
-    # combine list of dataframes into one df
-    print(f"Number of results: {len(result)}")
-
     for row in result:
         if row.empty:
             continue
@@ -167,11 +167,15 @@ def detect(
     # publish to debug topic
     if debug_publisher is not None:
         debug_publisher.publish(cv2_img.cv2_img_to_msg(cv_im))
+    if debug_inference_pub is not None:
         result = pd.concat(result)
-        result_paths = list(result["identity"])
-        result_images = [cv2.imread(path) for path in result_paths]
-        debug_publisher.publish(
-            cv2_img.cv2_img_to_msg(create_image_collage(result_images))
-        )
+        # check for empty result
+        if not result.empty:
+            result_paths = list(result["identity"])
+            result_images = [cv2.imread(path) for path in result_paths]
+            print(len(result_images))
+            debug_inference_pub.publish(
+                cv2_img.cv2_img_to_msg(create_image_collage(result_images))
+            )
 
     return response