Ros1 fixes for 3D Detection and 2D/3D tracking

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

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import os
 import rospy
 import time
@@ -27,6 +28,7 @@ def __init__(
         self,
         dataset: DatasetIterator,
         output_image_topic="/opendr/dataset_image",
+        data_fps=30,
     ):
         """
         Creates a ROS Node for publishing dataset images
@@ -36,6 +38,7 @@ def __init__(
         self.dataset = dataset
         # Initialize OpenDR ROSBridge object
         self.bridge = ROSBridge()
+        self.delay = 1.0 / data_fps
 
         if output_image_topic is not None:
             self.output_image_publisher = rospy.Publisher(
@@ -53,32 +56,60 @@ def start(self):
 
             rospy.loginfo("Publishing image [" + str(i) + "]")
             message = self.bridge.to_ros_image(
-                image, encoding="rgb8"
+                image, encoding="bgr8"
             )
             self.output_image_publisher.publish(message)
 
-            time.sleep(0.1)
+            time.sleep(self.delay)
 
             i += 1
 
 
-if __name__ == "__main__":
-
-    rospy.init_node('opendr_image_dataset')
-
-    dataset_path = MotDataset.download_nano_mot20(
-        "MOT", True
-    ).path
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-d", "--dataset_path", help="Path to a dataset",
+                        type=str, default="MOT")
+    parser.add_argument(
+        "-ks", "--mot20_subsets_path", help="Path to mot20 subsets",
+        type=str, default=os.path.join(
+            "..", "..", "src", "opendr", "perception", "object_tracking_2d",
+            "datasets", "splits", "nano_mot20.train"
+        )
+    )
+    parser.add_argument("-o", "--output_image_topic", help="Topic name to publish the data",
+                        type=str, default="/opendr/dataset_image")
+    parser.add_argument("-f", "--fps", help="Data FPS",
+                        type=float, default=30)
+    args = parser.parse_args()
+
+    dataset_path = args.dataset_path
+    mot20_subsets_path = args.mot20_subsets_path
+    output_image_topic = args.output_image_topic
+    data_fps = args.fps
+
+    if not os.path.exists(dataset_path):
+        dataset_path = MotDataset.download_nano_mot20(
+            "MOT", True
+        ).path
 
     dataset = RawMotDatasetIterator(
         dataset_path,
         {
-            "mot20": os.path.join(
-                "..", "..", "src", "opendr", "perception", "object_tracking_2d",
-                "datasets", "splits", "nano_mot20.train"
-            )
+            "mot20": mot20_subsets_path
         },
         scan_labels=False
     )
-    dataset_node = ImageDatasetNode(dataset)
+
+    rospy.init_node("image_dataset", anonymous=True)
+
+    dataset_node = ImageDatasetNode(
+        dataset,
+        output_image_topic=output_image_topic,
+        data_fps=data_fps,
+    )
+
     dataset_node.start()
+
+
+if __name__ == '__main__':
+    main()

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

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import torch
 import os
 import rospy
@@ -58,15 +59,13 @@ def __init__(
 
         self.learner.load(os.path.join(temp_dir, model_name), verbose=True)
 
-        # Initialize OpenDR ROSBridge object
+        self.input_point_cloud_topic = input_point_cloud_topic
         self.bridge = ROSBridge()
 
         self.detection_publisher = rospy.Publisher(
-            output_detection3d_topic, Detection3DArray, queue_size=10
+            output_detection3d_topic, Detection3DArray, queue_size=1
         )
 
-        rospy.Subscriber(input_point_cloud_topic, ROS_PointCloud, self.callback)
-
     def callback(self, data):
         """
         Callback that process the input data and publishes to the corresponding topics
@@ -80,39 +79,67 @@ def callback(self, data):
 
         # Convert detected boxes to ROS type and publish
         ros_boxes = self.bridge.to_ros_boxes_3d(detection_boxes, classes=["Car", "Van", "Truck", "Pedestrian", "Cyclist"])
-        if self.detection_publisher is not None:
-            self.detection_publisher.publish(ros_boxes)
-            rospy.loginfo("Published detection boxes")
+        self.detection_publisher.publish(ros_boxes)
+        rospy.loginfo("Published detection boxes")
+
+    def listen(self):
+        """
+        Start the node and begin processing input data.
+        """
+        rospy.init_node('voxel_detection_3d', anonymous=True)
+        rospy.Subscriber(self.input_point_cloud_topic, ROS_PointCloud, self.callback, queue_size=1, buff_size=10000000)
 
-if __name__ == "__main__":
-    # Automatically run on GPU/CPU
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        rospy.loginfo("Object Detection 3D Voxel Node started.")
+        rospy.spin()
 
-    # initialize ROS node
-    rospy.init_node("opendr_voxel_detection_3d", anonymous=True)
-    rospy.loginfo("Voxel Detection 3D node started")
 
-    model_name = rospy.get_param("~model_name", "tanet_car_xyres_16")
-    model_config_path = rospy.get_param(
-        "~model_config_path", os.path.join(
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--model_name", help="Name of the trained model",
+                        type=str, default="tanet_car_xyres_16")
+    parser.add_argument(
+        "-c", "--model_config_path", help="Path to a model .proto config",
+        type=str, default=os.path.join(
             "..", "..", "src", "opendr", "perception", "object_detection_3d",
             "voxel_object_detection_3d", "second_detector", "configs", "tanet",
-            "car", "test_short.proto"
+            "car", "xyres_16.proto"
         )
     )
-    temp_dir = rospy.get_param("~temp_dir", "temp")
-    input_point_cloud_topic = rospy.get_param(
-        "~input_point_cloud_topic", "/opendr/dataset_point_cloud"
-    )
-    rospy.loginfo("Using model_name: {}".format(model_name))
+    parser.add_argument("-t", "--temp_dir", help="Path to a temporary directory with models",
+                        type=str, default="temp")
+    parser.add_argument("-i", "--input_point_cloud_topic",
+                        help="Point Cloud topic provided by either a point_cloud_dataset_node or any other 3D Point Cloud Node",
+                        type=str, default="/opendr/dataset_point_cloud")
+    parser.add_argument("-o", "--output_detection3d_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/detection3d")
+    parser.add_argument("--device", help="Device to use, either \"cpu\" or \"cuda\", defaults to \"cuda\"",
+                        type=str, default="cuda", choices=["cuda", "cpu"])
+    args = parser.parse_args()
+    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"
 
-    # created node object
     voxel_node = ObjectDetection3DVoxelNode(
         device=device,
-        model_name=model_name,
-        model_config_path=model_config_path,
-        input_point_cloud_topic=input_point_cloud_topic,
-        temp_dir=temp_dir,
+        model_name=args.model_name,
+        model_config_path=args.model_config_path,
+        input_point_cloud_topic=args.input_point_cloud_topic,
+        temp_dir=args.temp_dir,
+        output_detection3d_topic=args.output_detection3d_topic,
     )
-    # begin ROS communications
-    rospy.spin()
+
+    voxel_node.listen()
+
+
+if __name__ == '__main__':
+    main()

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

@@ -13,10 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import cv2
 import torch
 import os
-from opendr.engine.target import TrackingAnnotation
+from opendr.engine.target import TrackingAnnotationList
 import rospy
 from vision_msgs.msg import Detection2DArray
 from std_msgs.msg import Int32MultiArray
@@ -35,9 +36,9 @@ def __init__(
         self,
         detector: Learner,
         input_image_topic="/usb_cam/image_raw",
-        output_detection_topic="/opendr/detection",
-        output_tracking_id_topic="/opendr/tracking_id",
-        output_image_topic="/opendr/image_annotated",
+        output_detection_topic="/opendr/deep_sort_detection",
+        output_tracking_id_topic="/opendr/deep_sort_tracking_id",
+        output_image_topic="/opendr/deep_sort_image_annotated",
         device="cuda:0",
         model_name="deep_sort",
         temp_dir="temp",
@@ -63,7 +64,6 @@ def __init__(
         :type temp_dir: str
         """
 
-        # # Initialize the face detector
         self.detector = detector
         self.learner = ObjectTracking2DDeepSortLearner(
             device=device, temp_path=temp_dir,
@@ -73,8 +73,9 @@ def __init__(
 
         self.learner.load(os.path.join(temp_dir, model_name), verbose=True)
 
-        # Initialize OpenDR ROSBridge object
         self.bridge = ROSBridge()
+        self.input_image_topic = input_image_topic
+
         self.tracking_id_publisher = rospy.Publisher(
             output_tracking_id_topic, Int32MultiArray, queue_size=10
         )
@@ -101,7 +102,6 @@ def callback(self, data):
         image = self.bridge.from_ros_image(data, encoding="bgr8")
         detection_boxes = self.detector.infer(image)
         image_with_detections = ImageWithDetections(image.numpy(), detection_boxes)
-        print(image_with_detections.data.shape)
         tracking_boxes = self.learner.infer(image_with_detections)
 
         if self.output_image_publisher is not None:
@@ -128,6 +128,16 @@ def callback(self, data):
             self.tracking_id_publisher.publish(ros_ids)
             rospy.loginfo("Published tracking ids")
 
+    def listen(self):
+        """
+        Start the node and begin processing input data.
+        """
+        rospy.init_node('object_tracking_2d_deep_sort_node', anonymous=True)
+        rospy.Subscriber(self.input_image_topic, ROS_Image, self.callback, queue_size=1, buff_size=10000000)
+
+        rospy.loginfo("Object Tracking 2D Deep Sort Node started.")
+        rospy.spin()
+
 
 colors = [
     (255, 0, 255),
@@ -139,7 +149,7 @@ def callback(self, data):
 ]
 
 
-def draw_predictions(frame, predictions: TrackingAnnotation, is_centered=False, is_flipped_xy=True):
+def draw_predictions(frame, predictions: TrackingAnnotationList, is_centered=False, is_flipped_xy=True):
     global colors
     w, h, _ = frame.shape
 
@@ -174,36 +184,62 @@ def draw_predictions(frame, predictions: TrackingAnnotation, is_centered=False,
         )
 
 
-if __name__ == "__main__":
-    # Automatically run on GPU/CPU
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-    # initialize ROS node
-    rospy.init_node("opendr_deep_sort", anonymous=True)
-    rospy.loginfo("Deep Sort node started")
-
-    model_name = rospy.get_param("~model_name", "deep_sort")
-    temp_dir = rospy.get_param("~temp_dir", "temp")
-    input_image_topic = rospy.get_param(
-        "~input_image_topic", "/opendr/dataset_image"
-    )
-    rospy.loginfo("Using model_name: {}".format(model_name))
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--model_name", help="Name of the trained model",
+                        type=str, default="deep_sort")
+    parser.add_argument("-t", "--temp_dir", help="Path to a temporary directory with models",
+                        type=str, default="temp")
+    parser.add_argument("-i", "--input_image_topic",
+                        help="Input Image topic provided by either an image_dataset_node, webcam or any other image node",
+                        type=str, default="/opendr/dataset_image")
+    parser.add_argument("-od", "--output_detection_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/deep_sort_detection")
+    parser.add_argument("-ot", "--output_tracking_id_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/deep_sort_tracking_id")
+    parser.add_argument("-oi", "--output_image_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/deep_sort_image_annotated")
+    parser.add_argument("--device", help="Device to use, either \"cpu\" or \"cuda\", defaults to \"cuda\"",
+                        type=str, default="cuda", choices=["cuda", "cpu"])
+    args = parser.parse_args()
+
+    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"
 
     detection_learner = ObjectTracking2DFairMotLearner(
-        device=device, temp_path=temp_dir,
+        device=device, temp_path=args.temp_dir,
     )
-    if not os.path.exists(os.path.join(temp_dir, "fairmot_dla34")):
-        ObjectTracking2DFairMotLearner.download("fairmot_dla34", temp_dir)
+    if not os.path.exists(os.path.join(args.temp_dir, "fairmot_dla34")):
+        ObjectTracking2DFairMotLearner.download("fairmot_dla34", args.temp_dir)
 
-    detection_learner.load(os.path.join(temp_dir, "fairmot_dla34"), verbose=True)
+    detection_learner.load(os.path.join(args.temp_dir, "fairmot_dla34"), verbose=True)
 
-    # created node object
     deep_sort_node = ObjectTracking2DDeepSortNode(
         detector=detection_learner,
         device=device,
-        model_name=model_name,
-        input_image_topic=input_image_topic,
-        temp_dir=temp_dir,
+        model_name=args.model_name,
+        input_image_topic=args.input_image_topic,
+        temp_dir=args.temp_dir,
+        output_detection_topic=args.output_detection_topic if args.output_detection_topic != "None" else None,
+        output_tracking_id_topic=args.output_tracking_id_topic if args.output_tracking_id_topic != "None" else None,
+        output_image_topic=args.output_image_topic if args.output_image_topic != "None" else None,
     )
-    # begin ROS communications
-    rospy.spin()
+
+    deep_sort_node.listen()
+
+
+if __name__ == '__main__':
+    main()