Qdtrack

configs/tracking.conf

@@ -4,7 +4,7 @@
 --obstacle_detection_gpu_memory_fraction=0.4
 ######### Tracker config #########
 --obstacle_tracking
---tracker_type=sort
+--tracker_type=qd_track
 #--tracker_type=deep_sort
 #--tracker_type=da_siam_rpn
 ###### Control config #####
@@ -15,7 +15,7 @@
 --simulator_num_people=200
 --simulator_num_vehicles=10
 ######### Other config #########
---visualize_tracked_obstacles
+# --visualize_tracked_obstacles
 --log_file_name=pylot.log
 --csv_log_file_name=pylot.csv
 --v=1

install.sh

@@ -79,6 +79,11 @@ mkdir -p tracking/center_track ; cd tracking/center_track
 # COCO model
 ~/.local/bin/gdown https://drive.google.com/uc?id=1tJCEJmdtYIh8VuN8CClGNws3YO7QGd40
 
+###### Download QDTrack models ######
+cd $PYLOT_HOME/dependencies/models
+mkdir -p tracking/qd_track ; cd tracking/qd_track
+~/.local/bin/gdown https://drive.google.com/uc?id=1YNAQgd8rMqqEG-fRj3VWlO4G5kdwJbxz
+
 ##### Download AnyNet depth estimation models #####
 echo "[x] Downloading the depth estimation models..."
 cd $PYLOT_HOME/dependencies/models
@@ -141,6 +146,15 @@ sudo apt-get install llvm-9
 export LLVM_CONFIG=/usr/bin/llvm-config-9
 python3 setup.py build develop --user
 
+###### Install QDTrack ######
+cd $PYLOT_HOME/dependencies/
+git clone https://github.com/mageofboy/qdtrack.git
+git clone https://github.com/open-mmlab/mmdetection.git
+cd mmdetection
+python3 setup.py develop #need to add mmcv
+cd $PYLOT_HOME/dependencies/qdtrack
+python3 setup.py develop
+
 ##### Download the Lanenet code #####
 echo "[x] Cloning the lanenet lane detection code..."
 cd $PYLOT_HOME/dependencies/

pylot/component_creator.py

@@ -351,6 +351,11 @@ def add_obstacle_tracking(center_camera_stream,
             obstacles_wo_history_tracking_stream = \
                 pylot.operator_creator.add_center_track_tracking(
                     center_camera_stream, center_camera_setup)
+        elif FLAGS.tracker_type == 'qd_track':
+            logger.debug('Using QDTrack obstacle tracker...')
+            obstacles_wo_history_tracking_stream = \
+                pylot.operator_creator.add_qd_track_tracking(
+                    center_camera_stream, center_camera_setup)
         else:
             logger.debug('Using obstacle tracker...')
             obstacles_wo_history_tracking_stream = \

pylot/flags.py

@@ -40,7 +40,7 @@
 flags.DEFINE_bool('perfect_obstacle_tracking', False,
                   'True to enable perfect obstacle tracking')
 flags.DEFINE_enum('tracker_type', 'sort',
-                  ['da_siam_rpn', 'deep_sort', 'sort', 'center_track'],
+                  ['da_siam_rpn', 'deep_sort', 'sort', 'center_track', 'qd_track'],
                   'Sets which obstacle tracker to use')
 flags.DEFINE_bool('lane_detection', False, 'True to enable lane detection')
 flags.DEFINE_bool('perfect_lane_detection', False,

pylot/operator_creator.py

@@ -269,6 +269,21 @@ def add_center_track_tracking(bgr_camera_stream,
     return obstacle_tracking_stream
 
 
+def add_qd_track_tracking(bgr_camera_stream,
+                              camera_setup,
+                              name='qd_track'):
+    from pylot.perception.tracking.qd_track_operator import \
+        QdTrackOperator
+    op_config = erdos.OperatorConfig(name='qd_track_operator',
+                                     log_file_name=FLAGS.log_file_name,
+                                     csv_log_file_name=FLAGS.csv_log_file_name,
+                                     profile_file_name=FLAGS.profile_file_name)
+    [obstacle_tracking_stream] = erdos.connect(QdTrackOperator, op_config,
+                                               [bgr_camera_stream], FLAGS,
+                                               camera_setup)
+    return obstacle_tracking_stream
+
+
 def add_tracking_evaluation(obstacle_tracking_stream,
                             ground_obstacles_stream,
                             evaluate_timely=False,

pylot/perception/flags.py

@@ -84,6 +84,16 @@
                   ['kitti_tracking', 'coco', 'mot', 'nuscenes'],
                   'CenterTrack available models')
 
+# QDTrack tracking flags.
+flags.DEFINE_string(
+    'qd_track_model_path',
+    'dependencies/models/tracking/qd_track/qdtrack-frcnn_r50_fpn_12e_bdd100k-13328aed.pth', #FIX
+    'Path to the model')
+flags.DEFINE_string(
+    'qd_track_config_path',
+    'dependencies/qdtrack/configs/qdtrack-frcnn_r50_fpn_12e_bdd100k.py', #FIX
+    'Path to the model')
+
 # Lane detection flags.
 flags.DEFINE_float('lane_detection_gpu_memory_fraction', 0.3,
                    'GPU memory fraction allocated to Lanenet')

pylot/perception/tracking/qd_track_operator.py

@@ -0,0 +1,86 @@
+import cv2
+import time
+
+import erdos
+
+import numpy as np
+
+from pylot.perception.detection.obstacle import Obstacle
+from pylot.perception.detection.utils import BoundingBox2D, \
+    BoundingBox3D, OBSTACLE_LABELS
+from pylot.perception.messages import ObstaclesMessage
+
+import torch
+
+class QdTrackOperator(erdos.Operator):
+    def __init__(self, camera_stream, obstacle_tracking_stream, flags,
+                 camera_setup):
+        from qdtrack.apis import init_model
+
+        camera_stream.add_callback(self.on_frame_msg,
+                                   [obstacle_tracking_stream])
+        self._flags = flags
+        self._logger = erdos.utils.setup_logging(self.config.name,
+                                                 self.config.log_file_name)
+        self._csv_logger = erdos.utils.setup_csv_logging(
+            self.config.name + '-csv', self.config.csv_log_file_name)
+        self._camera_setup = camera_setup
+        self.model = init_model(self._flags.qd_track_config_path, checkpoint=self._flags.qd_track_model_path, device='cuda:0', cfg_options=None)
+        self.classes = ('pedestrian', 'rider', 'car', 'bus', 'truck', 'bicycle', 'motorcycle', 'train')
+        self.frame_id = 0
+
+    @staticmethod
+    def connect(camera_stream):
+        obstacle_tracking_stream = erdos.WriteStream()
+        return [obstacle_tracking_stream]
+
+    def destroy(self):
+        self._logger.warn('destroying {}'.format(self.config.name))
+
+    @erdos.profile_method()
+    def on_frame_msg(self, msg, obstacle_tracking_stream):
+        """Invoked when a FrameMessage is received on the camera stream."""
+        from qdtrack.apis import inference_model
+
+        self._logger.debug('@{}: {} received frame'.format(
+            msg.timestamp, self.config.name))
+        assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
+        start_time = time.time()
+        image_np = msg.frame.as_bgr_numpy_array()
+        results = inference_model(self.model, image_np, self.frame_id)
+        self.frame_id += 1
+
+        bbox_result, track_result = results.values()
+        track_bboxes = np.zeros((0, 5))
+        track_ids = np.zeros((0), dtype=int)
+        track_labels = np.zeros((0), dtype=int)
+        obstacle_count = 0
+        for k, v in track_result.items():
+            track_bboxes = np.concatenate((track_bboxes, v['bbox'][None, :]), axis=0)
+            track_ids = np.concatenate((track_ids, np.array([k])), axis=0)
+            track_labels = np.concatenate((track_labels, np.array([v['label']])))
+            obstacle_count += 1
+        obstacles = []
+        for i in range(obstacle_count):
+            track_id = track_ids[i]
+            bbox = track_bboxes[i, :]
+            score = bbox[4]
+            label_id = track_labels[i]
+            label = self.classes[label_id]
+
+            if label in ['pedestrian', 'rider']:
+                label = 'person'
+
+            if label in OBSTACLE_LABELS:
+                bounding_box_2D = BoundingBox2D(bbox[0], bbox[2], bbox[1],
+                                                bbox[3])
+                bounding_box_3D = None
+                obstacles.append(
+                    Obstacle(bounding_box_2D,
+                             score,
+                             label,
+                             track_id,
+                             bounding_box_2D=bounding_box_2D))
+        runtime = (time.time() - start_time) * 1000
+        obstacle_tracking_stream.send(
+            ObstaclesMessage(msg.timestamp, obstacles, runtime))

requirements.txt

@@ -27,5 +27,7 @@ nuscenes-devkit
 progress
 pyquaternion
 scikit-learn==0.22.2
+mmcv>=0.3.0
+mmdet
 ##### CARLA dependencies #####
 networkx==2.2

scripts/lanenet_data_postprocess.py

@@ -0,0 +1,47 @@
+from absl import app, flags
+from pathlib import Path
+import numpy as np
+import os 
+from PIL import Image
+import matplotlib.pyplot as plt
+import math
+from random import sample
+FLAGS = flags.FLAGS
+flags.DEFINE_string('data_path', 'data/', 'binary images directory')
+
+def main(argv):
+    data_dir = os.path.abspath(FLAGS.data_path)
+    res = []
+    # ct = []
+    for f in os.listdir(data_dir):
+        curr_path = os.path.join(data_dir, f)
+        image = Image.open(curr_path)
+        im = np.asarray(image)
+        num_nonzero = np.count_nonzero(im)
+        # ct.append(num_nonzero)
+        if num_nonzero < 3000:
+            res.append(curr_path)
+    # n = math.ceil((max(ct) - min(ct))/1000)
+    # plt.hist(ct, bins=n)
+    # plt.show()
+    print(len(res))
+    res_sample = res
+    directories = ['gt_binary_image','gt_instance_image', 'gt_image']
+    filtered_dir = os.path.join(data_dir[:-len('gt_binary_image')], 'filtered/')
+    Path(filtered_dir).mkdir(parents=True, exist_ok=True)
+    for d in directories:
+        res_dir = os.path.join(filtered_dir, d)
+        Path(res_dir).mkdir(parents=True, exist_ok=True)
+    # Path(output_dir).mkdir(parents=True, exist_ok=True)
+    for f in res_sample:
+        s = f.split('gt_binary_image')
+        for d in directories:
+            curr_dir = d.join(s)
+            res_dir = os.path.join(filtered_dir, d) + s[1]
+            # print(curr_dir, res_dir)
+            Path(curr_dir).rename(res_dir)
+
+
+
+if __name__ == "__main__":
+    app.run(main)