Ros1 fixes heart anomaly detection

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

@@ -25,22 +25,23 @@
 
 class HeartAnomalyNode:
 
-    def __init__(self, input_topic="/ecg/ecg", prediction_topic="/opendr/heartanomaly", device="cuda", model='anbof'):
+    def __init__(self, input_ecg_topic="/ecg/ecg", output_heart_anomaly_topic="/opendr/heart_anomaly",
+                 device="cuda", model="anbof"):
         """
         Creates a ROS Node for heart anomaly (atrial fibrillation) detection from ecg data
-        :param input_topic: Topic from which we are reading the input array data
-        :type input_topic: str
-        :param prediction_topic: Topic to which we are publishing the predicted class
-        :type prediction_topic: str
+        :param input_ecg_topic: Topic from which we are reading the input array data
+        :type input_ecg_topic: str
+        :param output_heart_anomaly_topic: Topic to which we are publishing the predicted class
+        :type output_heart_anomaly_topic: str
         :param device: device on which we are running inference ('cpu' or 'cuda')
         :type device: str
         :param model: model to use: anbof or gru
         :type model: str
         """
 
-        self.publisher = rospy.Publisher(prediction_topic, Classification2D, queue_size=10)
+        self.publisher = rospy.Publisher(output_heart_anomaly_topic, Classification2D, queue_size=10)
 
-        rospy.Subscriber(input_topic, Float32MultiArray, self.callback)
+        rospy.Subscriber(input_ecg_topic, Float32MultiArray, self.callback)
 
         self.bridge = ROSBridge()
 
@@ -70,8 +71,8 @@ def listen(self):
     def callback(self, msg_data):
         """
         Callback that process the input data and publishes to the corresponding topics
-        :param data: input message
-        :type data: std_msgs.msg.Float32MultiArray
+        :param msg_data: input message
+        :type msg_data: std_msgs.msg.Float32MultiArray
         """
         # Convert Float32MultiArray to OpenDR Timeseries
         data = self.bridge.from_rosarray_to_timeseries(msg_data, self.channels, self.series_length)
@@ -83,17 +84,32 @@ def callback(self, msg_data):
         ros_class = self.bridge.from_category_to_rosclass(class_pred)
         self.publisher.publish(ros_class)
 
-if __name__ == '__main__':
-    # Select the device for running
-    try:
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    except:
-        device = 'cpu'
 
+if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('input_topic', type=str, help='listen to input data on this topic')
-    parser.add_argument('model', type=str, help='model to be used for prediction: anbof or gru')
+    parser.add_argument("--input_ecg_topic", type=str, default="/ecg/ecg",
+                        help="listen to input ECG data on this topic")
+    parser.add_argument("--model", type=str, default="anbof", help="model to be used for prediction: anbof or gru",
+                        choices=["anbof", "gru"])
+    parser.add_argument("--output_heart_anomaly_topic", type=str, default="/opendr/heart_anomaly",
+                        help="Topic name for heart anomaly detection topic")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to use (cpu, cuda)",
+                        choices=["cuda", "cpu"])
+
     args = parser.parse_args()
 
-    gesture_node = HeartAnomalyNode(input_topic=args.input_topic, model=args.model, device=device)
+    try:
+        if args.device == "cuda" and torch.cuda.is_available():
+            device = "cuda"
+        elif args.device == "cuda":
+            print("GPU not found. Using CPU instead.")
+            device = "cpu"
+        else:
+            print("Using CPU")
+            device = "cpu"
+    except:
+        print("Using CPU")
+        device = "cpu"
+
+    gesture_node = HeartAnomalyNode(input_ecg_topic=args.input_ecg_topic, model=args.model, device=device)
     gesture_node.listen()