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(
@@ -57,14 +60,32 @@ def start(self):
             )
             self.output_image_publisher.publish(message)
 
-            time.sleep(0.1)
+            time.sleep(self.delay)
 
             i += 1
 
 
-if __name__ == "__main__":
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-d", "--dataset_path", help="Path to a dataset",
+                        type=str, default="KITTI/opendr_nano_kitti")
+    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 upload the data",
+                        type=str, default="/opendr/dataset_image")
+    parser.add_argument("-f", "--fps", help="Data FPS",
+                        type=float, default=30)
+    args = parser.parse_args()
 
-    rospy.init_node('opendr_image_dataset')
+    dataset_path = args.dataset_path
+    mot20_subsets_path = args.mot20_subsets_path
+    output_image_topic = args.output_image_topic
+    data_fps = args.fps
 
     dataset_path = MotDataset.download_nano_mot20(
         "MOT", True
@@ -73,12 +94,17 @@ def start(self):
     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)
+    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
@@ -62,7 +63,7 @@ def __init__(
         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)
@@ -84,35 +85,53 @@ def callback(self, data):
             self.detection_publisher.publish(ros_boxes)
             rospy.loginfo("Published detection boxes")
 
-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_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 temp dir with models",
+                        type=str, default="temp")
+    parser.add_argument("-i", "--input_point_cloud_topic",
+                        help="Point Cloud topic provdied 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,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import cv2
 import torch
 import os
@@ -174,36 +175,61 @@ 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 temp dir with models",
+                        type=str, default="temp")
+    parser.add_argument("-i", "--input_image_topic",
+                        help="Input Image topic provdied 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/detection")
+    parser.add_argument("-ot", "--output_tracking_id_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/tracking_id")
+    parser.add_argument("-oi", "--output_image_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/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,
+        output_tracking_id_topic=args.output_tracking_id_topic,
+        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()

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

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import cv2
 import torch
 import os
@@ -166,27 +167,52 @@ 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"
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--model_name", help="Name of the trained model",
+                        type=str, default="fairmot_dla34")
+    parser.add_argument("-t", "--temp_dir", help="Path to a temp dir with models",
+                        type=str, default="temp")
+    parser.add_argument("-i", "--input_image_topic",
+                        help="Input Image topic provdied 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/detection")
+    parser.add_argument("-ot", "--output_tracking_id_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/tracking_id")
+    parser.add_argument("-oi", "--output_image_topic",
+                        help="Output detections topic",
+                        type=str, default="/opendr/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"
 
-    # initialize ROS node
-    rospy.init_node("opendr_fair_mot", anonymous=True)
-    rospy.loginfo("FairMOT node started")
-
-    model_name = rospy.get_param("~model_name", "fairmot_dla34")
-    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))
-
-    # created node object
     fair_mot_node = ObjectTracking2DFairMotNode(
         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,
+        output_tracking_id_topic=args.output_tracking_id_topic,
+        output_image_topic=args.output_image_topic if args.output_image_topic != "None" else None,
     )
-    # begin ROS communications
-    rospy.spin()
+
+    fair_mot_node.listen()
+
+if __name__ == '__main__':
+    main()