Configs + Code cleanup

README.md

@@ -86,18 +86,19 @@ export PYTHONPATH="${PYTHONPATH}:/juno/u/hkchiu/nuscenes-devkit/python-sdk"
 
 5. Choose a directory YOUR_NUSCENES_DETECTION_DIR (mine is /juno/u/hkchiu/dataset/nuscenes_new), and download then unzip the [MEGVII](https://github.com/poodarchu/Det3D)[3]'s detection results from the NuScenes [tracking webpage](https://www.nuscenes.org/data/detection-megvii.zip). 
 
-6. Edit our code: in main.py, evaluate_nuscenes.py, get_nuscenes_stats.py, replace “/juno/u/hkchiu/dataset/nuscenes_new/” with YOUR_NUSCENES_DETECTION_DIR, and replace “/juno/u/hkchiu/dataset/nuscenes/” with your YOUR_NUSCENES_DATA_DIR. 
+6. Edit configs: in `config/AB3DMOT.yaml` and `config/mahalanobis.yaml` set `DATA_ROOT` to `YOUR_NUSCENES_DATA_DIR`, and set `DETECTION_FILE` to your detection file in `YOUR_NUSCENES_DETECTION_DIR`. 
 
 7. Run the tracking code of our proposed method and the evaluation code (tracking takes 15 minutes, evaluation takes 90 minutes):
 ```
-python main.py val 2 m 11 greedy true nuscenes results/000008;
-python evaluate_nuscenes.py --output_dir results/000008 results/000008/val/results_val_probabilistic_tracking.json > results/000008/output.txt
+python main.py <CONFIG_FILE> <OUTPUT_DIR>
+python main.py config/mahalanobis.yaml results
+python evaluate_nuscenes.py --output_dir results/mahalanobis results/mahalanobis/megvii_val/results_val_probabilistic_tracking.json
 ```
 
 8. You can also run our implementation of the AB3DMOT[2] baseline method:
 ```
-python main.py val 0 iou 0.1 h false nuscenes results/000001;
-python evaluate_nuscenes.py --output_dir results/000001 results/000001/val/results_val_probabilistic_tracking.json > results/000001/output.txt
+python main.py config/AB3DMOT.yaml results
+python evaluate_nuscenes.py --output_dir results/AB3DMOT results/AB3DMOT/megvii_val/results_val_probabilistic_tracking.json
 ```
 
 9. Or you can run all the tracking methods (our implementation of AB3DMOT[2], our ablation methods, and our final proposed method) using the script:

algorithms.py

@@ -0,0 +1,83 @@
+# Metrics and matching algorithms for tracking
+
+import os.path, copy, numpy as np, time, sys
+from .utils.geometry_utils import diff_orientation_correction, convert_3dbox_to_8corner, iou3d
+import json
+
+from scipy.optimize import linear_sum_assignment as linear_assignment
+
+def pre_threshold_match(distance_matrix):
+  to_max_mask = distance_matrix > mahalanobis_threshold
+  distance_matrix[to_max_mask] = cfg.TRACKER.MATCH_THRESHOLD
+  matched_indices = linear_assignment(distance_matrix)      # hungarian algorithm
+  return np.transpose(np.asarray(matched_indices))
+
+def hungarian_match(distance_matrix):
+  return np.transpose(np.asarray(linear_assignment(distance_matrix)))                 # hungarian algorithm
+
+def greedy_match(distance_matrix):
+  '''
+  Find the one-to-one matching using greedy allgorithm choosing small distance
+  distance_matrix: (num_detections, num_tracks)
+  '''
+  matched_indices = []
+
+  num_detections, num_tracks = distance_matrix.shape
+  distance_1d = distance_matrix.reshape(-1)
+  index_1d = np.argsort(distance_1d)
+  index_2d = np.stack([index_1d // num_tracks, index_1d % num_tracks], axis=1)
+  detection_id_matches_to_tracking_id = [-1] * num_detections
+  tracking_id_matches_to_detection_id = [-1] * num_tracks
+  for sort_i in range(index_2d.shape[0]):
+    detection_id = int(index_2d[sort_i][0])
+    tracking_id = int(index_2d[sort_i][1])
+    if tracking_id_matches_to_detection_id[tracking_id] == -1 and detection_id_matches_to_tracking_id[detection_id] == -1:
+      tracking_id_matches_to_detection_id[tracking_id] = detection_id
+      detection_id_matches_to_tracking_id[detection_id] = tracking_id
+      matched_indices.append([detection_id, tracking_id])
+
+  matched_indices = np.array(matched_indices)
+  return matched_indices
+
+def mahalanobis_metric(detections, trackers, **kwargs):
+  """
+    Creates matrix of mahalanobis distances between detections and tracks
+
+    detections:  N x 7 
+    trackers:    M x 8
+    kwargs: {
+      trks_S: N x 7 x 7
+    }
+    Returns score matrix [M x N]
+    """
+  score_matrix = np.zeros((len(detections),len(trackers)),dtype=np.float32)
+  trks_S = kwargs['trks_S']
+  for d,det in enumerate(detections):
+    for t,trk in enumerate(trackers):
+      S_inv = np.linalg.inv(trks_S[t]) # 7 x 7
+      diff = np.expand_dims(det - trk, axis=1) # 7 x 1
+      # manual reversed angle by 180 when diff > 90 or < -90 degree
+      corrected_angle_diff = diff_orientation_correction(det[3], trk[3])
+      diff[3] = corrected_angle_diff
+      score_matrix[d, t] = np.sqrt(np.matmul(np.matmul(diff.T, S_inv), diff)[0][0])
+
+  return score_matrix
+
+def iou_metric(detections, trackers, **kwargs):
+  """
+    Creates matrix of negative IOU score between detections and tracks
+
+    detections:  N x 7 
+    trackers:    M x 8
+
+    Returns score matrix [M x N]
+    """
+  score_matrix = np.zeros((len(detections),len(trackers)),dtype=np.float32)
+
+  for d,det in enumerate(detections):
+    for t,trk in enumerate(trackers):
+      det_8corner = convert_3dbox_to_8corner(det, True)
+      trk_8corner = convert_3dbox_to_8corner(trk, True)
+      score_matrix[d,t] = -iou3d(det_8corner,trk_8corner)[0]   
+
+  return score_matrix

config/AB3DMOT.yaml

@@ -0,0 +1,24 @@
+
+DATASET:
+  NAME: NuScenes 
+  VERSION: 'v1.0-trainval'
+  DATA_ROOT: '/media/cyz/Data/cadc_tracker/PCDet/data/nuscenes'
+  DETECTION_FILE: '/media/cyz/Data/cadc_tracker/megvii_det/megvii_val.json'
+  DATA_SPLIT: 'val' 
+
+TRACKER:
+  COVARIANCE: 'BaselineCovariance'                # see covariance.py
+  MATCH_THRESHOLD: -0.1                           # -1 * Minimum IOU threshold 
+  MATCH_ALGORITHM: 'hungarian_match'              # see algorithms.py
+  SCORE_METRIC: 'iou_metric'
+  USE_ANGULAR_VELOCITY: False
+  CLASSES: [
+    'bicycle',
+    'bus',
+    'car',
+    'motorcycle',
+    'pedestrian',
+    'trailer',
+    'truck'
+  ]
+

config/mahalanobis.yaml

@@ -0,0 +1,24 @@
+
+DATASET:
+  NAME: NuScenes 
+  VERSION: 'v1.0-trainval'
+  DATA_ROOT: '/media/cyz/Data/cadc_tracker/PCDet/data/nuscenes'
+  DETECTION_FILE: '/media/cyz/Data/cadc_tracker/megvii_det/megvii_val.json'
+  DATA_SPLIT: 'val' 
+
+TRACKER:
+  COVARIANCE: 'NuScenesCovariance'       # see covariance.py
+  MATCH_THRESHOLD: 11                    # minimum mahalanobis distance
+  MATCH_ALGORITHM: 'greedy_match'        # see algorithms.py
+  SCORE_METRIC: 'mahalanobis_metric'
+  USE_ANGULAR_VELOCITY: True
+  CLASSES: [
+    'bicycle',
+    'bus',
+    'car',
+    'motorcycle',
+    'pedestrian',
+    'trailer',
+    'truck'
+  ]
+

covariance.py

@@ -1,106 +1,123 @@
 import numpy as np
+from .utils.config import cfg
 
-
-class Covariance(object):
+class BaselineCovariance(object):
   '''
-  Define different Kalman Filter covariance matrix
+  Define different Kalman Filter covariance matrix as in AB3DMOT
   Kalman Filter states:
   [x, y, z, rot_y, l, w, h, x_dot, y_dot, z_dot]
   '''
-  def __init__(self, covariance_id):
-    if covariance_id == 2:
-      self.num_states = 11 # with angular velocity
-    else:
-      self.num_states = 10
+  def __init__(self, **kwargs):
+    self.num_states = 10
     self.num_observations = 7
     self.P = np.eye(self.num_states)
     self.Q = np.eye(self.num_states)
     self.R = np.eye(self.num_observations)
 
-    NUSCENES_TRACKING_NAMES = [
-      'bicycle',
-      'bus',
-      'car',
-      'motorcycle',
-      'pedestrian',
-      'trailer',
-      'truck'
-    ]
+    self.P[self.num_observations:, self.num_observations:] *= 1000.
+    self.P *= 10.
+    self.Q[self.num_observations:, self.num_observations:] *= 0.01
+
+
+class KittiCovariance(object):
+  '''
+  Define different Kalman Filter covariance matrix based on Kitti data
+  Kalman Filter states:
+  [x, y, z, rot_y, l, w, h, x_dot, y_dot, z_dot]
+  '''
+  def __init__(self, **kwargs):
 
-    if covariance_id == 0:
-      # default from baseline code
-      self.P[self.num_observations:, self.num_observations:] *= 1000.
-      self.P *= 10.
-      self.Q[self.num_observations:, self.num_observations:] *= 0.01
-    elif covariance_id == 1:
-      # from kitti stats
-      self.P[0,0] = 0.01969623
-      self.P[1,1] = 0.01179107
-      self.P[2,2] = 0.04189842
-      self.P[3,3] = 0.52534431
-      self.P[4,4] = 0.11816206
-      self.P[5,5] = 0.00983173
-      self.P[6,6] = 0.01602004
-      self.P[7,7] = 0.01334779
-      self.P[8,8] = 0.00389245 
-      self.P[9,9] = 0.01837525
-
-      self.Q[0,0] = 2.94827444e-03
-      self.Q[1,1] = 2.18784125e-03
-      self.Q[2,2] = 6.85044585e-03
-      self.Q[3,3] = 1.10964054e-01
-      self.Q[4,4] = 0
-      self.Q[5,5] = 0
-      self.Q[6,6] = 0
-      self.Q[7,7] = 2.94827444e-03
-      self.Q[8,8] = 2.18784125e-03
-      self.Q[9,9] = 6.85044585e-03
+    self.num_states = 10
+    self.num_observations = 7
+    self.P = np.eye(self.num_states)
+    self.Q = np.eye(self.num_states)
+    self.R = np.eye(self.num_observations)
+
+    # from kitti stats
+    self.P[0,0] = 0.01969623
+    self.P[1,1] = 0.01179107
+    self.P[2,2] = 0.04189842
+    self.P[3,3] = 0.52534431
+    self.P[4,4] = 0.11816206
+    self.P[5,5] = 0.00983173
+    self.P[6,6] = 0.01602004
+    self.P[7,7] = 0.01334779
+    self.P[8,8] = 0.00389245 
+    self.P[9,9] = 0.01837525
 
-      self.R[0,0] = 0.01969623
-      self.R[1,1] = 0.01179107
-      self.R[2,2] = 0.04189842
-      self.R[3,3] = 0.52534431
-      self.R[4,4] = 0.11816206
-      self.R[5,5] = 0.00983173
-      self.R[6,6] = 0.01602004
+    self.Q[0,0] = 2.94827444e-03
+    self.Q[1,1] = 2.18784125e-03
+    self.Q[2,2] = 6.85044585e-03
+    self.Q[3,3] = 1.10964054e-01
+    self.Q[4,4] = 0
+    self.Q[5,5] = 0
+    self.Q[6,6] = 0
+    self.Q[7,7] = 2.94827444e-03
+    self.Q[8,8] = 2.18784125e-03
+    self.Q[9,9] = 6.85044585e-03
+
+    self.R[0,0] = 0.01969623
+    self.R[1,1] = 0.01179107
+    self.R[2,2] = 0.04189842
+    self.R[3,3] = 0.52534431
+    self.R[4,4] = 0.11816206
+    self.R[5,5] = 0.00983173
+    self.R[6,6] = 0.01602004
+
+
+class NuScenesCovariance(object):
+  '''
+  Define different Kalman Filter covariance matrix based on NuScenes data
+  This includes angular velocity
+  Kalman Filter states:
+  [x, y, z, rot_y, l, w, h, x_dot, y_dot, z_dot]
+  '''
+  def __init__(self, **kwargs):
 
-    elif covariance_id == 2:
-      # nuscenes
-      # see get_nuscenes_stats.py for the details on  how the numbers come from
-      #Kalman Filter state: [x, y, z, rot_z, l, w, h, x_dot, y_dot, z_dot, rot_z_dot]
+    self.num_states = 10
+    self.num_observations = 7
+    self.P = np.eye(self.num_states)
+    self.Q = np.eye(self.num_states)
+    self.R = np.eye(self.num_observations)
 
-      P = {
-        'bicycle':    [0.05390982, 0.05039431, 0.01863044, 1.29464435, 0.02713823, 0.01169572, 0.01295084, 0.04560422, 0.04097244, 0.01725477, 1.21635902],
-        'bus':        [0.17546469, 0.13818929, 0.05947248, 0.1979503 , 0.78867322, 0.05507407, 0.06684149, 0.13263319, 0.11508148, 0.05033665, 0.22529652],
-        'car':        [0.08900372, 0.09412005, 0.03265469, 1.00535696, 0.10912802, 0.02359175, 0.02455134, 0.08120681, 0.08224643, 0.02266425, 0.99492726],
-        'motorcycle': [0.04052819, 0.0398904 , 0.01511711, 1.06442726, 0.03291016, 0.00957574, 0.0111605 , 0.0437039 , 0.04327734, 0.01465631, 1.30414345],
-        'pedestrian': [0.03855275, 0.0377111 , 0.02482115, 2.0751833 , 0.02286483, 0.0136347 , 0.0203149 , 0.04237008, 0.04092393, 0.01482923, 2.0059979 ],
-        'trailer':    [0.23228021, 0.22229261, 0.07006275, 1.05163481, 1.37451601, 0.06354783, 0.10500918, 0.2138643 , 0.19625241, 0.05231335, 0.97082174],
-        'truck':      [0.14862173, 0.1444596 , 0.05417157, 0.73122169, 0.69387238, 0.05484365, 0.07748085, 0.10683797, 0.10248689, 0.0378078 , 0.76188901]
-      }
+    # nuscenes
+    # see get_nuscenes_stats.py for the details on  how the numbers come from
+    #Kalman Filter state: [x, y, z, rot_z, l, w, h, x_dot, y_dot, z_dot, rot_z_dot]
 
-      Q = {
-        'bicycle':    [1.98881347e-02, 1.36552276e-02, 5.10175742e-03, 1.33430252e-01, 0, 0, 0, 1.98881347e-02, 1.36552276e-02, 5.10175742e-03, 1.33430252e-01],
-        'bus':        [1.17729925e-01, 8.84659079e-02, 1.17616440e-02, 2.09050032e-01, 0, 0, 0, 1.17729925e-01, 8.84659079e-02, 1.17616440e-02, 2.09050032e-01],
-        'car':        [1.58918523e-01, 1.24935318e-01, 5.35573165e-03, 9.22800791e-02, 0, 0, 0, 1.58918523e-01, 1.24935318e-01, 5.35573165e-03, 9.22800791e-02],
-        'motorcycle': [3.23647590e-02, 3.86650974e-02, 5.47421635e-03, 2.34967407e-01, 0, 0, 0, 3.23647590e-02, 3.86650974e-02, 5.47421635e-03, 2.34967407e-01],
-        'pedestrian': [3.34814566e-02, 2.47354921e-02, 5.94592529e-03, 4.24962535e-01, 0, 0, 0, 3.34814566e-02, 2.47354921e-02, 5.94592529e-03, 4.24962535e-01],
-        'trailer':    [4.19985099e-02, 3.68661552e-02, 1.19415050e-02, 5.63166240e-02, 0, 0, 0, 4.19985099e-02, 3.68661552e-02, 1.19415050e-02, 5.63166240e-02],
-        'truck':      [9.45275998e-02, 9.45620374e-02, 8.38061721e-03, 1.41680460e-01, 0, 0, 0, 9.45275998e-02, 9.45620374e-02, 8.38061721e-03, 1.41680460e-01]
-      }
+    P = {
+      'bicycle':    [0.05390982, 0.05039431, 0.01863044, 1.29464435, 0.02713823, 0.01169572, 0.01295084, 0.04560422, 0.04097244, 0.01725477, 1.21635902],
+      'bus':        [0.17546469, 0.13818929, 0.05947248, 0.1979503 , 0.78867322, 0.05507407, 0.06684149, 0.13263319, 0.11508148, 0.05033665, 0.22529652],
+      'car':        [0.08900372, 0.09412005, 0.03265469, 1.00535696, 0.10912802, 0.02359175, 0.02455134, 0.08120681, 0.08224643, 0.02266425, 0.99492726],
+      'motorcycle': [0.04052819, 0.0398904 , 0.01511711, 1.06442726, 0.03291016, 0.00957574, 0.0111605 , 0.0437039 , 0.04327734, 0.01465631, 1.30414345],
+      'pedestrian': [0.03855275, 0.0377111 , 0.02482115, 2.0751833 , 0.02286483, 0.0136347 , 0.0203149 , 0.04237008, 0.04092393, 0.01482923, 2.0059979 ],
+      'trailer':    [0.23228021, 0.22229261, 0.07006275, 1.05163481, 1.37451601, 0.06354783, 0.10500918, 0.2138643 , 0.19625241, 0.05231335, 0.97082174],
+      'truck':      [0.14862173, 0.1444596 , 0.05417157, 0.73122169, 0.69387238, 0.05484365, 0.07748085, 0.10683797, 0.10248689, 0.0378078 , 0.76188901]
+    }
 
-      R = {
-        'bicycle':    [0.05390982, 0.05039431, 0.01863044, 1.29464435, 0.02713823, 0.01169572, 0.01295084],
-        'bus':        [0.17546469, 0.13818929, 0.05947248, 0.1979503 , 0.78867322, 0.05507407, 0.06684149],
-        'car':        [0.08900372, 0.09412005, 0.03265469, 1.00535696, 0.10912802, 0.02359175, 0.02455134],
-        'motorcycle': [0.04052819, 0.0398904 , 0.01511711, 1.06442726, 0.03291016, 0.00957574, 0.0111605 ],
-        'pedestrian': [0.03855275, 0.0377111 , 0.02482115, 2.0751833 , 0.02286483, 0.0136347 , 0.0203149 ],
-        'trailer':    [0.23228021, 0.22229261, 0.07006275, 1.05163481, 1.37451601, 0.06354783, 0.10500918],
-        'truck':      [0.14862173, 0.1444596 , 0.05417157, 0.73122169, 0.69387238, 0.05484365, 0.07748085]
-      }
-
-      self.P = {tracking_name: np.diag(P[tracking_name]) for tracking_name in NUSCENES_TRACKING_NAMES}
-      self.Q = {tracking_name: np.diag(Q[tracking_name]) for tracking_name in NUSCENES_TRACKING_NAMES}
-      self.R = {tracking_name: np.diag(R[tracking_name]) for tracking_name in NUSCENES_TRACKING_NAMES}
-    else:
-      assert(False)
+    Q = {
+      'bicycle':    [1.98881347e-02, 1.36552276e-02, 5.10175742e-03, 1.33430252e-01, 0, 0, 0, 1.98881347e-02, 1.36552276e-02, 5.10175742e-03, 1.33430252e-01],
+      'bus':        [1.17729925e-01, 8.84659079e-02, 1.17616440e-02, 2.09050032e-01, 0, 0, 0, 1.17729925e-01, 8.84659079e-02, 1.17616440e-02, 2.09050032e-01],
+      'car':        [1.58918523e-01, 1.24935318e-01, 5.35573165e-03, 9.22800791e-02, 0, 0, 0, 1.58918523e-01, 1.24935318e-01, 5.35573165e-03, 9.22800791e-02],
+      'motorcycle': [3.23647590e-02, 3.86650974e-02, 5.47421635e-03, 2.34967407e-01, 0, 0, 0, 3.23647590e-02, 3.86650974e-02, 5.47421635e-03, 2.34967407e-01],
+      'pedestrian': [3.34814566e-02, 2.47354921e-02, 5.94592529e-03, 4.24962535e-01, 0, 0, 0, 3.34814566e-02, 2.47354921e-02, 5.94592529e-03, 4.24962535e-01],
+      'trailer':    [4.19985099e-02, 3.68661552e-02, 1.19415050e-02, 5.63166240e-02, 0, 0, 0, 4.19985099e-02, 3.68661552e-02, 1.19415050e-02, 5.63166240e-02],
+      'truck':      [9.45275998e-02, 9.45620374e-02, 8.38061721e-03, 1.41680460e-01, 0, 0, 0, 9.45275998e-02, 9.45620374e-02, 8.38061721e-03, 1.41680460e-01]
+    }
+
+    R = {
+      'bicycle':    [0.05390982, 0.05039431, 0.01863044, 1.29464435, 0.02713823, 0.01169572, 0.01295084],
+      'bus':        [0.17546469, 0.13818929, 0.05947248, 0.1979503 , 0.78867322, 0.05507407, 0.06684149],
+      'car':        [0.08900372, 0.09412005, 0.03265469, 1.00535696, 0.10912802, 0.02359175, 0.02455134],
+      'motorcycle': [0.04052819, 0.0398904 , 0.01511711, 1.06442726, 0.03291016, 0.00957574, 0.0111605 ],
+      'pedestrian': [0.03855275, 0.0377111 , 0.02482115, 2.0751833 , 0.02286483, 0.0136347 , 0.0203149 ],
+      'trailer':    [0.23228021, 0.22229261, 0.07006275, 1.05163481, 1.37451601, 0.06354783, 0.10500918],
+      'truck':      [0.14862173, 0.1444596 , 0.05417157, 0.73122169, 0.69387238, 0.05484365, 0.07748085]
+    }
+
+    self.P = np.diag(P[kwargs['tracking_name']])
+    self.Q = np.diag(Q[kwargs['tracking_name']])
+    self.R = np.diag(R[kwargs['tracking_name']])
+
+    if not cfg.TRACKER.USE_ANGULAR_VELOCITY:
+        self.P = self.P[:-1,:-1]
+        self.Q = self.Q[:-1,:-1]

evaluate_nuscenes.py

@@ -24,8 +24,7 @@
                         help='Folder to store result metrics, graphs and example visualizations.')
     parser.add_argument('--eval_set', type=str, default='val',
                         help='Which dataset split to evaluate on, train, val or test.')
-    #parser.add_argument('--dataroot', type=str, default='/data/sets/nuscenes',
-    parser.add_argument('--dataroot', type=str, default='/juno/u/hkchiu/dataset/nuscenes/trainval',
+    parser.add_argument('--dataroot', type=str, default='/media/cyz/Data/cadc_tracker/PCDet/data/nuscenes',
                         help='Default nuScenes data directory.')
     parser.add_argument('--version', type=str, default='v1.0-trainval',
                         help='Which version of the nuScenes dataset to evaluate on, e.g. v1.0-trainval.')