rewrite happypose node into a service

scripts/happypose_with_camera.py

@@ -9,7 +9,10 @@ from happypose.toolbox.inference.types import (
 )
 
 import rospy
+import tf
+from ros_cosypose.srv import Inference, InferenceResponse
 from sensor_msgs.msg import Image, CameraInfo
+from std_msgs.msg import Int32MultiArray
 from geometry_msgs.msg import Pose, Point, Quaternion
 from pytorch3d import transforms
 
@@ -24,33 +27,54 @@ def run_inference_from_realtime_camera(
         observation = ObservationTensor.from_numpy(image.data, None, camera_info.K)
         CosyPose = CosyPoseWrapper(dataset_name=dataset_to_use, n_workers=8)
         predictions = CosyPose.inference(observation)
-        return predictions.poses
+        rospy.loginfo()
+        return predictions.poses, predictions.infos
     except AttributeError as err:
-        print("rien trouvé")
-        return []
+        rospy.loginfo("rien trouvé")
+        return [], []
 
 
-if __name__ == "__main__":
+def retrieve_parameters_and_run_inference():
     dataset = rospy.get_param("~dataset", "ycbv")
     model = rospy.get_param("~model", "megapose-1.0-RGB-multi-hypothesis")
-
-    rospy.init_node("caemra_listener", anonymous=True)
     cam_info = rospy.wait_for_message("/camera/color/camera_info", CameraInfo)
-
     image = rospy.wait_for_message("/camera/color/image_raw", Image)
-    inference = run_inference_from_realtime_camera(cam_info, image, model, dataset)
-    for i, pose in enumerate(inference):
+    return run_inference_from_realtime_camera(cam_info, image, model, dataset)
+
+
+def convert_list_to_ros_msg(list):
+    res = Int32MultiArray()
+    res.data.append(list)
+    return res.data
+
+
+def get_nb_objects_in_image(req):
+    inference, inference_infos = retrieve_parameters_and_run_inference()
+    nb_objects_in_image = [0 for _ in range(31)]
+    if inference_infos == []:
+        return convert_list_to_ros_msg(nb_objects_in_image)
+    labels = inference_infos["label"]
+    for idx, pose in enumerate(inference):
+        object_number = labels[idx][-2:]
+        if not object_number.isnumeric():
+            raise RuntimeError(
+                "label doesn't end with tless object index : ", object_number
+            )
+        object_number = int(object_number)
+        nb_objects_in_image[object_number] += 1
         quat = transforms.matrix_to_quaternion(pose[:3, :3])
-        pose = Pose(
-            position=Point(pose[0, 3], pose[1, 3], pose[2, 3]),
-            orientation=Quaternion(quat[0], quat[1], quat[2], quat[3]),
-        )
-        pub = rospy.Publisher(f"tf/happypose_object_{i}", Pose, queue_size=10)
-        pub.publish(pose)
-    if inference == []:
-        pose = Pose(
-            position=Point(0, 0, 0),
-            orientation=Quaternion(1, 0, 0, 0),
+        br = tf.TransformBroadcaster()
+        br.sendTransform(
+            (pose[0, 3], pose[1, 3], pose[2, 3]),
+            (quat[0], quat[1], quat[2], quat[3]),
+            rospy.Time.now(),
+            f"tf/happypose_object_{object_number}_{nb_objects_in_image[object_number]}",
+            "camera_color_optical_frame",
         )
-        pub = rospy.Publisher("tf/happypose_object", Pose, queue_size=10)
-        pub.publish(pose)
+    return convert_list_to_ros_msg(nb_objects_in_image)
+
+
+if __name__ == "__main__":
+    rospy.init_node("inference_on_camera_image", anonymous=True)
+    s = rospy.Service("happypose_inference", Inference, get_nb_objects_in_image)
+    rospy.spin()