diff trafficintelligence/moving.py @ 1253:ef68d4ba7dae

added loading ego vehicle in kitti 2D format and method to plot outline
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Mon, 25 Mar 2024 17:05:20 -0400
parents fe35473acee3
children 20a5e1292321
line wrap: on
line diff
--- a/trafficintelligence/moving.py	Fri Mar 22 14:33:25 2024 -0400
+++ b/trafficintelligence/moving.py	Mon Mar 25 17:05:20 2024 -0400
@@ -1414,6 +1414,7 @@
                  5: 'bus',
                  #6: 'train',
                  7: 'truck'}
+kitti2Types = {'Car': 1, 'Cyclist': 4, 'Pedestrian': 2, 'Van': 1}
 
 userType2Num = utils.inverseEnumeration(userTypeNames)
 
@@ -1440,6 +1441,8 @@
         self.userType = userType
         self.setNObjects(nObjects) # a feature has None for nObjects
         self.features = features
+        if features is not None and len(features) > 1:
+            self.featureNumbers = [f.getNum() for f in features]
         # compute bounding polygon from trajectory
 
     @staticmethod
@@ -1819,15 +1822,21 @@
             for f in self.getFeatures():
                 f.setStationary()
         
-    def plot(self, options = '', withOrigin = False, timeStep = 1, withFeatures = False, withIds = False, **kwargs):
+    def plot(self, options = '', withOrigin = False, timeStep = 1, withFeatures = False, withOutline = False, withIds = False, **kwargs):
         if withIds:
             objNum = self.getNum()
         else:
             objNum = None
-        if withFeatures and self.hasFeatures():
-            for f in self.getFeatures():
-                f.positions.plot('r', True, timeStep, **kwargs)
-            self.positions.plot('bx-', True, timeStep, objNum, **kwargs)
+        if (withFeatures or withOutline) and self.hasFeatures():
+            if withOutline:
+                for t in self.getTimeInterval():
+                    if t%timeStep == 0:
+                        self.plotOutlineAtInstant(t, cvutils.colors[self.getNum()%len(cvutils.colors)]+'-')
+                self.positions.plot(cvutils.colors[self.getNum()%len(cvutils.colors)]+'-', withOrigin, timeStep, objNum, **kwargs)
+            else:
+                for f in self.getFeatures():
+                    f.positions.plot('r', True, timeStep, **kwargs)
+                self.positions.plot('bx-', True, timeStep, objNum, **kwargs)
         else:
             self.positions.plot(options, withOrigin, timeStep, objNum, **kwargs)