Updated, cleaned up and added argparse to ROS1 fall detection node

projects/opendr_ws/src/perception/scripts/fall_detection.py

@@ -13,76 +13,89 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import cv2
+import argparse
+import torch
 
 import rospy
-import torch
-import cv2
 from vision_msgs.msg import Detection2DArray
 from sensor_msgs.msg import Image as ROS_Image
 from opendr_bridge import ROSBridge
+
+from opendr.engine.data import Image
+from opendr.engine.target import BoundingBox, BoundingBoxList
 from opendr.perception.pose_estimation import get_bbox
 from opendr.perception.pose_estimation import LightweightOpenPoseLearner
 from opendr.perception.fall_detection import FallDetectorLearner
-from opendr.engine.data import Image
-from opendr.engine.target import BoundingBox, BoundingBoxList
 
 
 class FallDetectionNode:
 
-    def __init__(self, input_image_topic="/usb_cam/image_raw", output_image_topic="/opendr/image_fall_annotated",
-                 fall_annotations_topic="/opendr/falls", device="cuda"):
+    def __init__(self, input_rgb_image_topic="/usb_cam/image_raw",
+                 output_rgb_image_topic="/opendr/image_fallen_annotated", detections_topic="/opendr/fallen",
+                 device="cuda", num_refinement_stages=2, use_stride=False, half_precision=False):
         """
-        Creates a ROS Node for fall detection
-        :param input_image_topic: Topic from which we are reading the input image
-        :type input_image_topic: str
-        :param output_image_topic: Topic to which we are publishing the annotated image (if None, we are not publishing
-        annotated image)
-        :type output_image_topic: str
-        :param fall_annotations_topic: Topic to which we are publishing the annotations (if None, we are not publishing
-        annotated fall annotations)
-        :type fall_annotations_topic:  str
+        Creates a ROS Node for rule-based fall detection based on Lightweight OpenPose.
+        :param input_rgb_image_topic: Topic from which we are reading the input image
+        :type input_rgb_image_topic: str
+        :param output_rgb_image_topic: Topic to which we are publishing the annotated image (if None, no annotated
+        image is published)
+        :type output_rgb_image_topic: str
+        :param detections_topic: Topic to which we are publishing the annotations (if None, no fall detection message
+        is published)
+        :type detections_topic:  str
         :param device: device on which we are running inference ('cpu' or 'cuda')
         :type device: str
+        :param num_refinement_stages: Specifies the number of pose estimation refinement stages are added on the
+        model's head, including the initial stage. Can be 0, 1 or 2, with more stages meaning slower and more accurate
+        inference
+        :type num_refinement_stages: int
+        :param use_stride: Whether to add a stride value in the model, which reduces accuracy but increases
+        inference speed
+        :type use_stride: bool
+        :param half_precision: Enables inference using half (fp16) precision instead of single (fp32) precision.
+        Valid only for GPU-based inference
+        :type half_precision: bool
         """
-        if output_image_topic is not None:
-            self.image_publisher = rospy.Publisher(output_image_topic, ROS_Image, queue_size=10)
+        self.input_rgb_image_topic = input_rgb_image_topic
+
+        if output_rgb_image_topic is not None:
+            self.image_publisher = rospy.Publisher(output_rgb_image_topic, ROS_Image, queue_size=1)
         else:
             self.image_publisher = None
 
-        if fall_annotations_topic is not None:
-            self.fall_publisher = rospy.Publisher(fall_annotations_topic, Detection2DArray, queue_size=10)
+        if detections_topic is not None:
+            self.fall_publisher = rospy.Publisher(detections_topic, Detection2DArray, queue_size=1)
         else:
             self.fall_publisher = None
 
-        self.input_image_topic = input_image_topic
-
         self.bridge = ROSBridge()
 
-        # Initialize the pose estimation
-        self.pose_estimator = LightweightOpenPoseLearner(device=device, num_refinement_stages=2,
-                                                         mobilenet_use_stride=False,
-                                                         half_precision=False)
+        # Initialize the pose estimation learner
+        self.pose_estimator = LightweightOpenPoseLearner(device=device, num_refinement_stages=num_refinement_stages,
+                                                         mobilenet_use_stride=use_stride,
+                                                         half_precision=half_precision)
         self.pose_estimator.download(path=".", verbose=True)
         self.pose_estimator.load("openpose_default")
 
+        # Initialize the fall detection learner
         self.fall_detector = FallDetectorLearner(self.pose_estimator)
 
     def listen(self):
         """
-        Start the node and begin processing input data
+        Start the node and begin processing input data.
         """
-        rospy.init_node('opendr_fall_detection', anonymous=True)
-        rospy.Subscriber(self.input_image_topic, ROS_Image, self.callback)
-        rospy.loginfo("Fall detection node started!")
+        rospy.init_node('fall_detection_node', anonymous=True)
+        rospy.Subscriber(self.input_rgb_image_topic, ROS_Image, self.callback, queue_size=1, buff_size=5000000)
+        rospy.loginfo("Fall detection node started.")
         rospy.spin()
 
     def callback(self, data):
         """
-        Callback that process the input data and publishes to the corresponding topics
-        :param data: input message
+        Callback that processes the input data and publishes to the corresponding topics.
+        :param data: Input image message
         :type data: sensor_msgs.msg.Image
         """
-
         # Convert sensor_msgs.msg.Image into OpenDR Image
         image = self.bridge.from_ros_image(data, encoding='bgr8')
 
@@ -100,34 +113,64 @@ def callback(self, data):
             if fallen == 1:
                 color = (0, 0, 255)
                 x, y, w, h = get_bbox(pose)
-                bbox = BoundingBox(left=x, top=y, width=w, height=h, name=0)
-                bboxes.data.append(bbox)
 
                 cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
-                cv2.putText(image, "Detected fallen person", (5, 55), cv2.FONT_HERSHEY_SIMPLEX,
-                            0.75, color, 1, cv2.LINE_AA)
+                cv2.putText(image, "Fallen person", (x, y + h - 10), cv2.FONT_HERSHEY_SIMPLEX,
+                            1, color, 2, cv2.LINE_AA)
 
                 # Convert detected boxes to ROS type and publish
+                bboxes.data.append(BoundingBox(left=x, top=y, width=w, height=h, name="fallen person"))
                 ros_boxes = self.bridge.to_ros_boxes(bboxes)
                 if self.fall_publisher is not None:
                     self.fall_publisher.publish(ros_boxes)
 
         if self.image_publisher is not None:
-            message = self.bridge.to_ros_image(Image(image), encoding='bgr8')
-            self.image_publisher.publish(message)
-
+            self.image_publisher.publish(self.bridge.to_ros_image(Image(image), encoding='bgr8'))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_rgb_image_topic", help="Topic name for input rgb image",
+                        type=str, default="/usb_cam/image_raw")
+    parser.add_argument("-o", "--output_rgb_image_topic", help="Topic name for output annotated rgb image",
+                        type=str, default="/opendr/image_fallen_annotated")
+    parser.add_argument("-d", "--detections_topic", help="Topic name for detection messages",
+                        type=str, default="/opendr/fallen")
+    parser.add_argument("--device", help="Device to use, either \"cpu\" or \"cuda\", defaults to \"cuda\"",
+                        type=str, default="cuda", choices=["cuda", "cpu"])
+    parser.add_argument("--accelerate", help="Enables acceleration flags (e.g., stride)", default=False,
+                        action="store_true")
+    args = parser.parse_args()
 
-if __name__ == '__main__':
-    # Select the device for running the
     try:
-        if torch.cuda.is_available():
-            print("GPU found.")
-            device = 'cuda'
-        else:
+        if args.device == "cuda" and torch.cuda.is_available():
+            device = "cuda"
+        elif args.device == "cuda":
             print("GPU not found. Using CPU instead.")
-            device = 'cpu'
+            device = "cpu"
+        else:
+            print("Using CPU.")
+            device = "cpu"
     except:
-        device = 'cpu'
-
-    fall_detection_node = FallDetectionNode(device=device)
+        print("Using CPU.")
+        device = "cpu"
+
+    if args.accelerate:
+        stride = True
+        stages = 0
+        half_prec = True
+    else:
+        stride = False
+        stages = 2
+        half_prec = False
+
+    fall_detection_node = FallDetectionNode(device=device,
+                                            input_rgb_image_topic=args.input_rgb_image_topic,
+                                            output_rgb_image_topic=args.output_rgb_image_topic,
+                                            detections_topic=args.detections_topic,
+                                            num_refinement_stages=stages, use_stride=stride, half_precision=half_prec)
     fall_detection_node.listen()
+
+
+if __name__ == '__main__':
+    main()