msight-base 0.1.2__py3-none-any.whl

msight_base/__init__.py

@@ -0,0 +1,3 @@
+from .detection import *
+from .road_user import *
+from .trajectory import *

msight_base/behavior.py

@@ -0,0 +1,36 @@
+from enum import Enum
+
+
+class BehaviorType(Enum):
+    """
+    Enum for different types of behaviors.
+    """
+    UNKNOWN = -1
+    CONSTANT = 0
+    ACCELERATE = 1
+    DECELERATE = 2
+    STOP = 3
+    LANE_KEEPING = 4
+    LANE_CHANGING = 5
+    LANE_DEPARTURE = 6
+    LEFT_TURN = 7
+    RIGHT_TURN = 8
+
+    # Additional behaviors
+    U_TURN = 9
+    YIELDING = 10
+    FOLLOWING = 11
+    OVERTAKING = 12
+    PARKING = 13
+    EMERGENCY_BRAKING = 14
+
+    def __str__(self):
+        return self.name.lower()
+    
+    @classmethod
+    def from_name(cls, name: str):
+        try:
+            return cls[name.upper()]
+        except KeyError:
+            # Fall back to UNKNOWN
+            return cls.UNKNOWN

msight_base/detection.py

@@ -0,0 +1,42 @@
+from typing import List, Optional
+from .utils.cls import get_class_path, import_class_from_path
+
+class DetectedObjectBase:
+    def __init__(self):
+        pass
+
+    def to_dict(self):
+        raise NotImplementedError("Subclasses should implement this!")
+
+class DetectionResultBase:
+    def __init__(self, object_list: List[DetectedObjectBase], timestamp: int, sensor_type: Optional[str] = None):
+        self.object_list = object_list
+        self.timestamp = timestamp
+        self.sensor_type = sensor_type
+
+    def to_dict(self):
+        object_data_type = get_class_path(self.object_list[0].__class__) if len(self.object_list) > 0 else get_class_path(DetectedObjectBase)
+        return {
+            'object_list': [obj.to_dict() for obj in self.object_list],
+            'timestamp': self.timestamp,
+            'sensor_type': self.sensor_type,
+            'object_data_type': object_data_type,
+        }
+    
+    @staticmethod
+    def from_dict(data: dict) -> 'DetectionResultBase':
+        object_data_type = data.get("object_data_type")
+        if object_data_type is None:
+            raise ValueError(f"Missing 'object_data_type' in data: {data}")
+
+        obj_cls = import_class_from_path(object_data_type)
+        if data['object_list'] is not None and len(data['object_list']) > 0:
+            object_list = [obj_cls.from_dict(obj_data) for obj_data in data['object_list']]
+        else:
+            object_list = []
+
+        return DetectionResultBase(
+            object_list=object_list,
+            timestamp=data['timestamp'],
+            sensor_type=data['sensor_type'],
+        )

msight_base/map.py

@@ -0,0 +1,213 @@
+from enum import Enum
+from typing import List
+from commonroad.scenario.lanelet import LaneletType, LaneletNetwork
+import matplotlib.pyplot as plt
+
+
+class LaneShape(Enum):
+    """
+    Enum for different types of lane shape.
+    """
+    STRAIGHT = 'straight'
+    LEFT_TURN = 'left-turn'
+    RIGHT_TURN = 'right-turn'
+    U_TURN = 'u-turn'
+
+
+class LaneSide(Enum):
+    """
+    Enum for different types of lane side.
+    """
+    LEFT = 1
+    RIGHT = -1
+    ONLANE = 0
+    UNKNOWN = 'unknown'
+
+    @classmethod
+    def from_name(cls, name: str):
+        try:
+            return cls[name.upper()]
+        except KeyError:
+            # Fall back to UNKNOWN
+            return cls.UNKNOWN
+
+
+class MapInfo:
+    def __init__(self, lane_id, lane_point_idx, dis_to_lane_center, side: LaneSide, related_lane_id, related_route, nearest_next_lane_point_idx=None):
+        self.lane_id = lane_id
+        self.lane_point_idx = lane_point_idx
+        self.dis_to_lane_center = dis_to_lane_center
+        self.side = side
+        self.related_lane_id = related_lane_id
+        self.related_route = related_route
+        self.nearest_next_lane_point_idx = nearest_next_lane_point_idx
+
+    def __repr__(self):
+        return f"MapInfo(lane_id={self.lane_id}, lane_point_idx={self.lane_point_idx})"
+    
+    def from_dict(cls, object_dict):
+        return cls(
+            lane_id=object_dict.get('lane_id', None),
+            lane_point_idx=object_dict.get('lane_point_idx', None),
+            dis_to_lane_center=object_dict.get('dis_to_lane_center', None),
+            side=LaneSide.from_name(object_dict.get('side', 'unknown')),
+            related_lane_id=object_dict.get('related_lane_id', None),
+            related_route=object_dict.get('related_route', None),
+            nearest_next_lane_point_idx=object_dict.get('nearest_next_lane_point_idx', None)
+        )
+    
+    def to_dict(self):
+        return {
+            'lane_id': self.lane_id,
+            'lane_point_idx': self.lane_point_idx,
+            'dis_to_lane_center': self.dis_to_lane_center,
+            'side': self.side.value,
+            'related_lane_id': self.related_lane_id,
+            'related_route': self.related_route,
+            'nearest_next_lane_point_idx': self.nearest_next_lane_point_idx
+        }
+
+
+class MapObject:
+    def __init__(self, commonroadObj: LaneletNetwork, center_point, map_lanes, lane_heading, lane_shape: List[LaneShape], lane_seg_length, background_img, corner_coords):
+        self.center_point = center_point
+        self._map_lanes = map_lanes
+        self._lane_heading = lane_heading
+        self._lane_shape = lane_shape
+        self._lane_seg_length = lane_seg_length
+        self.background_img = background_img
+        self.corner_coords = corner_coords
+
+        if isinstance(commonroadObj, LaneletNetwork):
+            self.map_polylines_ids = [lanelet.lanelet_id for lanelet in commonroadObj.lanelets]
+            self._suc_edges = [lanelet.successor for lanelet in commonroadObj.lanelets]
+            self._pre_edges = [lanelet.predecessor for lanelet in commonroadObj.lanelets]
+            self._left_edges = [lanelet.adj_left for lanelet in commonroadObj.lanelets]
+            self._left_edge_directions = [lanelet.adj_left_same_direction for lanelet in commonroadObj.lanelets]
+            self._right_edges = [lanelet.adj_right for lanelet in commonroadObj.lanelets]
+            self._right_edge_directions = [lanelet.adj_right_same_direction for lanelet in commonroadObj.lanelets]
+            self.intersection_lane_id_list = [lanelet.lanelet_id for lanelet in commonroadObj.lanelets if LaneletType.INTERSECTION in lanelet.lanelet_type]
+            self.straight_lane_id_list = [lanelet.lanelet_id for lanelet in commonroadObj.lanelets if LaneletType.INTERSECTION not in lanelet.lanelet_type]
+        else:
+            self.map_polylines_ids = []
+            self._suc_edges = []
+            self._pre_edges = []
+            self._left_edges = []
+            self._left_edge_directions = []
+            self._right_edges = []
+            self._right_edge_directions = []
+            self.intersection_lane_id_list = []
+            self.straight_lane_id_list = []
+
+    def suc_edges(self, id=-1):
+        if id == -1:
+            return self._suc_edges
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._suc_edges[self.map_polylines_ids.index(id)]
+
+    def pre_edges(self, id=-1):
+        if id == -1:
+            return self._pre_edges
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._pre_edges[self.map_polylines_ids.index(id)]
+
+    def left_edges(self, id=-1):
+        if id == -1:
+            return self._left_edges
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._left_edges[self.map_polylines_ids.index(id)]
+
+    def right_edges(self, id=-1):
+        if id == -1:
+            return self._right_edges
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._right_edges[self.map_polylines_ids.index(id)]
+
+    def left_edge_directions(self, id=-1):
+        if id == -1:
+            return self._left_edge_directions
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._left_edge_directions[self.map_polylines_ids.index(id)]
+
+    def right_edge_directions(self, id=-1):
+        if id == -1:
+            return self._right_edge_directions
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._right_edge_directions[self.map_polylines_ids.index(id)]
+
+    def lane_shape(self, id=-1):
+        if id == -1:
+            return self._lane_shape
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            return self._lane_shape[self.map_polylines_ids.index(id)]
+
+    def map_lanes(self, id=-1, idx=-1):
+        if id == -1:
+            return self._map_lanes
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            if idx == -1:
+                return self._map_lanes[self.map_polylines_ids.index(id)]
+            else:
+                if idx >= len(self._map_lanes[self.map_polylines_ids.index(id)]):
+                    raise ValueError(f"Index {idx} out of range for lane with id {id}")
+                return self._map_lanes[self.map_polylines_ids.index(id)][idx]
+
+    def lane_heading(self, id=-1, idx=-1):
+        if id == -1:
+            return self._lane_heading
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            if idx == -1:
+                return self._lane_heading[self.map_polylines_ids.index(id)]
+            else:
+                if idx >= len(self._lane_heading[self.map_polylines_ids.index(id)]):
+                    raise ValueError(f"Index {idx} out of range for lane with id {id}")
+                return self._lane_heading[self.map_polylines_ids.index(id)][idx]
+    
+    def lane_seg_length(self, id=-1, idx=-1):
+        if id == -1:
+            return self._lane_seg_length
+        else:
+            if id not in self.map_polylines_ids:
+                raise ValueError(f"Lanelet id {id} not found in map_polylines_ids")
+            if idx == -1:
+                return self._lane_seg_length[self.map_polylines_ids.index(id)]
+            else:
+                if idx >= len(self._lane_seg_length[self.map_polylines_ids.index(id)]):
+                    raise ValueError(f"Index {idx} out of range for lane with id {id}")
+                return self._lane_seg_length[self.map_polylines_ids.index(id)][idx]
+
+    def to_dict(self):
+        return {
+            'center_point': self.center_point,
+            'map_lanes': self._map_lanes,
+            'lane_heading': self._lane_heading,
+            'map_polylines_ids': self.map_polylines_ids,
+            'suc_edges': self._suc_edges,
+            'pre_edges': self._pre_edges,
+            'left_edges': self._left_edges,
+            'left_edge_directions': self._left_edge_directions,
+            'right_edges': self._right_edges,
+            'right_edge_directions': self._right_edge_directions,
+            'lane_shape': self._lane_shape,
+            'intersection_lane_id_list': self.intersection_lane_id_list,
+            'straight_lane_id_list': self.straight_lane_id_list,
+            'corner_coords': self.corner_coords,
+        }

msight_base/road_user.py

@@ -0,0 +1,182 @@
+from enum import IntEnum
+from msight_base.map import MapInfo
+from msight_base.behavior import BehaviorType
+
+
+class RoadUserPoint:
+    def __init__(self,
+                 x=None,
+                 y=None,
+                 speed=None,
+                 acceleration=None,
+                 heading=None,
+                 width=None,
+                 length=None,
+                 height=None,
+                 poly_box=None,
+                 category=None,
+                 confidence=None,
+                 turning_signal=None,
+                 map_info=None,
+                 timestamp=None,
+                 frame_step=None,
+                 traj_id=None,
+                 sensor_data={},
+                 behaviors=[],
+                 conf_int_2sigma=None,
+                 conf_int_vel_2sigma=None,
+                 heading_confidence=None,
+                 yaw_rate=None):
+        self._timestamp = timestamp
+        self._frame_step = frame_step
+        self.x = x
+        self.y = y
+        self.speed = speed
+        self.acceleration = acceleration
+        self.heading = heading
+        self.width = width
+        self.length = length
+        self.height = height
+        self.poly_box = poly_box
+        self.category = category
+        self.traj = None
+        self.frame = None
+        self.prev = None
+        self.next = None
+        self.sensor_data = sensor_data # sensor id to sensor data map
+        self.confidence = confidence
+        self.turning_signal = turning_signal
+        self.map_info = map_info
+        self.behaviors = behaviors
+        self.pred_trajectory = None
+        self.conf_int_2sigma = conf_int_2sigma
+        self.conf_int_vel_2sigma = conf_int_vel_2sigma
+        self.heading_confidence = heading_confidence
+        self.yaw_rate = yaw_rate
+
+        ## this is used to store the id when the object's trajectory is not yet created or assigned
+        self._traj_id = traj_id
+
+        ## this is also used to store the id when the object's trajectory is not yet created or assigned
+        ## for some use cases, we need a uuid instead of a trajectory id (that increments from 0)
+        self._traj_uuid = None
+
+    @property
+    def traj_id(self):
+        if self.traj is not None:
+            return self.traj.id
+        return self._traj_id
+
+    @traj_id.setter
+    def traj_id(self, value):
+        if self.traj is not None:
+            raise AttributeError("Cannot set traj_id directly when traj is assigned.")
+        self._traj_id = value
+
+    @property
+    def frame_step(self):
+        if self.frame is None:
+            return None
+        return self.frame.step
+    
+    @frame_step.setter
+    def frame_step(self, value):
+        if self.frame is not None:
+            raise AttributeError("Cannot set frame_step directly when frame is assigned.")
+        self._frame_step = value
+
+    @property
+    def timestamp(self):
+        if self.frame is None:
+            return self._timestamp
+        return self.frame.timestamp
+    
+    @timestamp.setter
+    def timestamp(self, value):
+        if self.frame is not None:
+            raise AttributeError("Cannot set timestamp directly when frame is assigned.")
+        self._timestamp = value
+
+    @staticmethod
+    def from_dict(object_dict):
+        return RoadUserPoint(
+            timestamp=object_dict.get('timestamp', None),
+            frame_step=object_dict.get('frame_step', None),
+            traj_id=object_dict.get('traj_id', None),
+            x=object_dict.get('x', None),
+            y=object_dict.get('y', None),
+            speed=object_dict.get('speed', None),
+            acceleration=object_dict.get('acceleration', None),
+            heading=object_dict.get('heading', None),
+            width=object_dict.get('width', None),
+            length=object_dict.get('length', None),
+            height=object_dict.get('height', None),
+            poly_box=object_dict.get('poly_box', None),
+            category=RoadUserCategory(int(object_dict['category'])) if object_dict.get('category') is not None else None,
+            confidence=object_dict.get('confidence', None),
+            turning_signal=object_dict.get('turning_signal', None),
+            map_info=MapInfo.from_dict(object_dict['map_info']) if object_dict.get('map_info') else None,
+            sensor_data=object_dict.get('sensor_data', {}),
+            behaviors=[BehaviorType[behavior] for behavior in object_dict.get('behaviors', [])],
+            conf_int_2sigma=object_dict.get('conf_int_2sigma', None),
+            conf_int_vel_2sigma=object_dict.get('conf_int_vel_2sigma', None),
+            heading_confidence=object_dict.get('heading_confidence', None),
+            yaw_rate=object_dict.get('yaw_rate', None)
+        )
+
+    def to_dict(self):
+        return {
+            'traj_id': self.traj_id,
+            'frame_step': self.frame_step,
+            'timestamp': self.timestamp,
+            'x': self.x,
+            'y': self.y,
+            'speed': self.speed,
+            'acceleration': self.acceleration,
+            'heading': self.heading,
+            'width': self.width,
+            'length': self.length,
+            'height': self.height,
+            'category': RoadUserCategory(int(self.category)) if self.category is not None else None,
+            'confidence': self.confidence,
+            'turning_signal': self.turning_signal,
+            'poly_box': self.poly_box,
+            'map_info': self.map_info.to_dict() if self.map_info else None,
+            'behaviors': [str(behavior) for behavior in self.behaviors],
+            'sensor_data': self.sensor_data,
+            'conf_int_2sigma': self.conf_int_2sigma,
+            'conf_int_vel_2sigma': self.conf_int_vel_2sigma,
+            'heading_confidence': self.heading_confidence,
+            'yaw_rate': self.yaw_rate,
+        }
+
+    def __repr__(self):
+        return f"RoadUserPoint(x={self.x}, y={self.y}, heading={self.heading}, width={self.width}, length={self.length}, category={self.category})"
+
+
+class RoadUserCategory(IntEnum):
+    """
+    Enum for road user categories.
+    """
+    UNKNOWN = -1
+    SEDAN = 0
+    SUV = 1
+    TRUCK = 2
+    BUS = 3
+    PEDESTRIAN = 4
+    BICYCLE = 5
+    MOTORCYCLE = 6
+    TRAILER = 7
+    VAN = 8
+    PICKUP = 9
+
+    def __str__(self):
+        return self.name.lower()
+    
+    @classmethod
+    def from_name(cls, name: str):
+        try:
+            return cls[name.upper()]
+        except KeyError:
+            # Fall back to UNKNOWN
+            return cls.UNKNOWN

msight_base/trajectory.py

@@ -0,0 +1,186 @@
+import bisect
+from .road_user import RoadUserPoint
+from typing import List
+
+class Container:
+    def __init__(self, id=None, objects = None):
+        self.id = id
+        if objects is None:
+            objects = []
+        self.objects = objects
+
+    def __len__(self):
+        return len(self.objects)
+
+    def __iter__(self):
+        return iter(self.objects)
+    
+    def __getitem__(self, item):
+        return self.objects[item]
+
+
+class Trajectory(Container):
+    def __init__(self, id=None):
+        self.steps = []
+        self.step_to_object_map = {}
+        super().__init__(id=id, objects=[])
+
+    def add_object(self, obj, step, insort=False):
+        if insort:
+            index = bisect.bisect_left(self.steps, step)
+            if index < len(self.steps) and self.steps[index] == step:
+                raise ValueError(f"Step {step} already exists")
+            self.steps.insert(index, step)
+            self.objects.insert(index, obj)
+            self.step_to_object_map[step] = obj
+            if index > 0:
+                self.objects[index - 1].next = obj
+                obj.prev = self.objects[index - 1]
+            if index < len(self.objects) - 1:
+                obj.next = self.objects[index + 1]
+                self.objects[index + 1].prev = obj
+        else:
+            if len(self.steps) > 0 and step < self.steps[-1]:
+                raise ValueError(f"Step {step} is less than the last step {self.steps[-1]}, if you want to insert in between use insort=True")
+            self.objects.append(obj)
+            self.step_to_object_map[step] = obj
+            self.steps.append(step)
+            if len(self.objects) > 1:
+                self.objects[-2].next = obj
+                obj.prev = self.objects[-2]
+        obj.traj = self
+
+    def get_object_at_step(self, step):
+        return self.step_to_object_map.get(step, None)
+    
+    def remove_object(self, step):
+        if step not in self.steps:
+            return
+            # raise ValueError(f"Step {step} does not exist in the trajectory")
+        obj = self.get_object_at_step(step)
+        self.objects.remove(obj)
+        self.steps.remove(step)
+        del self.step_to_object_map[step]
+        obj.traj = None
+        next_obj = obj.next
+        prev_obj = obj.prev
+        if next_obj is not None:
+            next_obj.prev = prev_obj
+        if prev_obj is not None:
+            prev_obj.next = next_obj
+
+
+class Frame(Container):
+    def __init__(self, step, timestamp=None):
+        self.step = step
+        self.timestamp = timestamp
+        self.traj_ids = set()
+        self.traj_id_to_obj_map = {}
+        super().__init__(None, objects=[])
+
+    def add_object(self, obj):
+        if obj.traj_id in self.traj_ids:
+            raise ValueError(f"Object with id {obj.traj_id} already exists in the frame")
+        self.objects.append(obj)
+        self.traj_id_to_obj_map[obj.traj_id] = obj
+        self.traj_ids.add(obj.traj_id)
+        obj.frame = self
+
+    def remove_object(self, obj):
+        if obj.traj_id not in self.traj_ids:
+            raise ValueError(f"Object with id {obj.traj_id} does not exist in the frame")
+        self.objects.remove(obj)
+        del self.traj_id_to_obj_map[obj.traj_id]
+        self.traj_ids.remove(obj.traj_id)
+        obj.frame = None
+
+
+class TrajectoryManager:
+    def __init__(self, max_frames=None):
+        self.trajectories = []
+        self.traj_ids = set()
+        self.traj_id_to_traj_map = {}
+        self.frames = []
+        self.steps = []
+        self.timestamps = []
+        self.step_to_frame_map = {}
+        self.timestamp_to_frame_map = {}
+        self.max_frames = max_frames
+
+    @property
+    def last_step(self):
+        return self.steps[-1] if len(self.steps)>0 else -1
+    
+    @property
+    def earliest_step(self):
+        return self.steps[0] if len(self.steps)>0 else -1
+
+    @property
+    def last_frame(self):
+        return self.frames[-1]
+    
+    def delete_earliest_frame(self):
+        if not self.frames:
+            raise ValueError("No frames to delete")
+        step = self.frames[0].step
+        for obj in self.frames[0].objects:
+            # print(obj.traj, obj.traj_id)
+            if obj.traj is not None:
+                obj.traj.remove_object(step)
+        if self.frames[0].timestamp is not None:
+            del self.timestamp_to_frame_map[self.frames[0].timestamp]
+            self.timestamps.remove(self.frames[0].timestamp)
+        del self.step_to_frame_map[step]
+        self.steps.remove(step)
+
+        for traj in self.trajectories:
+            if len(traj.objects) == 0:
+                self.remove_traj(traj)
+        self.frames.pop(0)
+
+    def add_object(self, obj, traj_id, step, timestamp=None, insort=False):
+        if traj_id not in self.traj_ids:
+            traj = Trajectory(traj_id)
+            self.trajectories.append(traj)
+            self.traj_ids.add(traj_id)
+            self.traj_id_to_traj_map[traj_id] = traj
+        else:
+            traj = self.traj_id_to_traj_map[traj_id]
+
+        if step >= self.last_step + 1:
+            frame = Frame(step, timestamp)
+            self.frames.append(frame)
+            self.steps.append(step)
+            self.step_to_frame_map[step] = frame
+            if timestamp is not None:
+                self.timestamps.append(timestamp)
+                self.timestamp_to_frame_map[timestamp] = frame
+            
+        elif step <= self.last_step and step >= self.earliest_step:
+            frame = self.step_to_frame_map[step]
+        else:
+            raise ValueError(f"Step {step} is not a valid step, valid steps are from 0 to {self.last_step+1}")
+
+        traj.add_object(obj, step, insort=insort)
+        frame.add_object(obj)
+        while self.max_frames is not None and len(self.frames) > self.max_frames:
+            # print(len(self.frames), self.max_frames)
+            self.delete_earliest_frame()
+
+    def add_list_as_new_frame(self, object_list: List[RoadUserPoint], timestamp=None):
+        # use this method to push a list of objects as the last frame, this is very useful when you have a list of objects that are already tracked with id assigned
+        step = self.last_step + 1
+        for obj in object_list:
+            if obj.traj_id is None:
+                raise ValueError("Object must have a traj_id to be added to a trajectory manager")
+            self.add_object(obj, obj.traj_id, step, timestamp=obj.timestamp)
+
+
+    def remove_traj(self, traj: Trajectory):
+        # print(f"Removing trajectory with id {traj.id}")
+        tid = traj.id
+        self.trajectories.remove(traj)
+        self.traj_ids.remove(tid)
+        del self.traj_id_to_traj_map[tid]
+        for obj in traj:
+            obj.frame.remove_object(obj)

msight_base/utils/cls.py

@@ -0,0 +1,9 @@
+def get_class_path(cls):
+    module = cls.__module__
+    class_name = cls.__name__
+    return f"{module}.{class_name}"
+
+def import_class_from_path(path: str):
+    module_path, class_name = path.rsplit('.', 1)
+    module = __import__(module_path, fromlist=[class_name])
+    return getattr(module, class_name)

msight_base/utils/data.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+from msight_base import TrajectoryManager, RoadUserPoint
+import json
+from datetime import datetime
+from tqdm import tqdm
+
+
+def get_timestamp_from_filename(file_name: str) -> datetime:
+    # file name is like "2023-09-05 10-00-00-039784.json"
+    stem = file_name.split(".")[0]
+    return datetime.strptime(stem, "%Y-%m-%d %H-%M-%S-%f")
+
+
+def read_msight_json_data(file_path: Path) -> TrajectoryManager:
+    tm = TrajectoryManager()
+    step = 0
+    for frame_file in tqdm(list(file_path.iterdir())):
+        if frame_file.suffix != ".json":
+            continue
+        timestamp = get_timestamp_from_filename(frame_file.stem)
+        with open(frame_file, "r") as f:
+            frame_data = json.load(f)
+            for obj_data in frame_data['fusion']:
+                obj = RoadUserPoint(
+                    obj_data['lat'],
+                    obj_data['lon'],
+                    heading=obj_data['heading'],
+                    width=obj_data['width'],
+                    length=obj_data['length'],
+                )
+                tm.add_object(obj, obj_data['id'], step, timestamp=timestamp)
+        step += 1
+    return tm

msight_base/visualizer/__init__.py

@@ -0,0 +1 @@
+from .visualizer import Visualizer

msight_base/visualizer/utils.py

@@ -0,0 +1,118 @@
+import numpy as np
+import cv2
+import json
+
+
+class VideoWriter(object):
+
+    def __init__(self, name, out_size_w=960, out_size_h=720):
+
+        self.video_writer = cv2.VideoWriter(
+            name + '.mp4', cv2.VideoWriter_fourcc(*'MP4V'),
+            10.0, (out_size_w, out_size_h))
+        self.video_writer_cat = cv2.VideoWriter(
+            name + '_cat.mp4', cv2.VideoWriter_fourcc(*'MP4V'),
+            10.0, (out_size_h + out_size_w, out_size_h))
+
+        self.out_size_w = out_size_w
+        self.out_size_h = out_size_h
+
+    def save_frame(self, vis_det, vis_loc):
+
+        vis_det = cv2.resize(vis_det, (self.out_size_w, self.out_size_h))
+        vis_loc = cv2.resize(vis_loc, (self.out_size_h, self.out_size_h))
+
+        self.video_writer.write(vis_det[:, :, ::-1])
+
+        frame_cat = np.concatenate([vis_det, vis_loc], axis=1)
+        self.video_writer_cat.write(frame_cat[:, :, ::-1])
+
+        cv2.namedWindow('frame_cat', cv2.WINDOW_NORMAL)
+        cv2.imshow('frame_cat', frame_cat[:, :, ::-1])
+        cv2.waitKey(1)
+
+
+############ some helper functions ##############
+
+class Struct:
+    def __init__(self, **entries):
+        self.__dict__.update(entries)
+
+
+def parse_config(path_to_json=r'./config.json'):
+    with open(path_to_json) as f:
+      data = json.load(f)
+    return Struct(**data)
+
+
+
+################ BB draw on 2D image... ####################
+
+def draw_bb_on_image(vehicle_list, img):
+
+    for i in range(len(vehicle_list)):
+        v = vehicle_list[i]
+        x_bottom_c, y_bottom_c = v.pixel_bottom_center.x, v.pixel_bottom_center.y
+        diagonal_length = v.diagonal_length_pixel
+
+        # draw box
+        pt1 = (int(x_bottom_c - diagonal_length / 2), int(y_bottom_c - diagonal_length / 2))
+        pt2 = (int(x_bottom_c + diagonal_length / 2), int(y_bottom_c + diagonal_length / 2))
+        cv2.rectangle(img, pt1=pt1, pt2=pt2, color=(255, 0, 0), thickness=1, lineType=cv2.LINE_AA)
+        box_tmp = img[pt1[1]:pt2[1], pt1[0]:pt2[0], :]
+        img[pt1[1]:pt2[1], pt1[0]:pt2[0], 1:] = (0.75 * box_tmp[:, :, 1:]).astype(np.uint8)
+
+        # draw bottom center
+        cv2.circle(img, (x_bottom_c, y_bottom_c), 2, (255, 0, 0), -1)
+        # text = 'id=%s' % v.id
+        # pt = (x_bottom_c, y_bottom_c - 5)
+        # cv2.putText(img, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        #             fontScale=0.5, color=(255, 255, 0), thickness=1, lineType=cv2.LINE_AA)
+
+    return img
+
+
+################ BB (YOLOX) draw on 2D image... ####################
+# Instead of draw the square, we directly draw the rect predicted by YOLOX
+
+def draw_bb_on_image_yolox(vehicle_list, img):
+
+    for i in range(len(vehicle_list)):
+        v = vehicle_list[i]
+        x_bottom_c, y_bottom_c = v.pixel_bottom_center.x, v.pixel_bottom_center.y
+        x1, y1, x2, y2 = v.pixel_bbox
+
+        # draw box
+        pt1 = (int(x1), int(y1))
+        pt2 = (int(x2), int(y2))
+        cv2.rectangle(img, pt1=pt1, pt2=pt2, color=(255, 0, 0), thickness=1, lineType=cv2.LINE_AA)
+        box_tmp = img[pt1[1]:pt2[1], pt1[0]:pt2[0], :]
+        img[pt1[1]:pt2[1], pt1[0]:pt2[0], 1:] = (0.75 * box_tmp[:, :, 1:]).astype(np.uint8)
+
+        # draw bottom center
+        cv2.circle(img, (x_bottom_c, y_bottom_c), 2, (255, 0, 0), -1)
+        # text = 'id=%s' % v.id
+        # pt = (x_bottom_c, y_bottom_c - 5)
+        # cv2.putText(img, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        #             fontScale=0.5, color=(255, 255, 0), thickness=1, lineType=cv2.LINE_AA)
+
+    return img
+
+# map center lat/lon
+# use the following for Ann Arbor Ellsworth Roundabout
+LAT0, LON0 = 42.229392, -83.739012
+
+
+def coord_normalization(lat, lon, center_lat=LAT0, center_lon=LON0):
+    "from lat/lon to local coordinate with unit meter"
+    lat_norm = (lat - center_lat) * 111000.
+    lon_norm = (lon - center_lon) * 111000. * np.cos(center_lat/180.*np.pi)
+    return lat_norm, lon_norm
+
+
+def coord_unnormalization(lat_norm, lon_norm, center_lat=LAT0, center_lon=LON0):
+    "from local coordinate with unit meter to lat/lon"
+    lat = lat_norm / 111000. + center_lat
+    lon = lon_norm / 111000. / np.cos(center_lat/180.*np.pi) + center_lon
+    return lat, lon
+

msight_base/visualizer/visualizer.py

@@ -0,0 +1,289 @@
+import numpy as np
+import cv2
+import os
+import json
+import math
+from msight_base import TrajectoryManager
+from .utils import coord_normalization, coord_unnormalization
+
+
+class Struct:
+    def __init__(self, **entries):
+        self.__dict__.update(entries)
+
+
+def parse_config(path_to_json=r'./config.json'):
+    with open(path_to_json) as f:
+        data = json.load(f)
+    return Struct(**data)
+
+
+class Visualizer:
+    """
+    Basemap for detection visualization. Show and map detection to base map layer.
+    Pixel image: vehicle bounding box, vehicle id
+    Map layer: location, heading, rect box, trajectory, predicted future...
+    """
+
+    def __init__(self, map_image_path, map_width=1024, map_height=1024, RGB=False):
+        self.h, self.w = map_height, map_width
+
+        self.f = parse_config(os.path.splitext(map_image_path)[0] + '.json')
+        self.transform_wd2px, self.transform_px2wd = self._create_coord_mapper()
+
+        basemap = cv2.imread(map_image_path, cv2.IMREAD_COLOR)
+        if not RGB:
+            basemap = cv2.cvtColor(basemap, cv2.COLOR_BGR2RGB)
+        self.basemap = cv2.resize(basemap, (map_width, map_height))
+        self.basemap = (self.basemap.astype(np.float64) * 0.3).astype(np.uint8)
+
+        np.random.seed(0)
+        self.color_table = np.random.randint(80, 255, (10, 3))
+        np.random.seed()
+
+        self.traj_manager = TrajectoryManager(max_frames=100)
+        self.traj_layer = np.zeros([self.h, self.w, 3], dtype=np.uint8)
+        self.traj_alpha = np.zeros([self.h, self.w, 3], dtype=np.float32)
+
+        if os.path.exists(r'./vehicle_category.json'):
+            self.label_list = parse_config(r'./vehicle_category.json')
+        else:
+            self.label_list = None
+    
+    @staticmethod
+    def _draw_vehicle_as_point(vis, ptc, color):
+        # draw circle boundary
+        cv2.circle(vis, tuple(ptc), 10, color, -1)
+        # draw center location
+        cv2.circle(vis, tuple(ptc), 4, (255, 255, 0), -1)
+
+    @staticmethod
+    def _draw_vehicle_as_box(vis, pts, color):
+        # fill rectangle
+        cv2.fillPoly(vis, [np.array(pts)], color, lineType=cv2.LINE_AA)
+
+    @staticmethod
+    def _draw_vehicle_heading_as_arrow(vis, ptc, heading, color):
+        # draw heading
+        # convert degree to radian
+        heading = math.radians(heading-90)
+
+        line_length = 25
+        pt1 = (int(ptc[0]), int(ptc[1]))
+        pt2 = (int(ptc[0] + line_length*np.cos(heading)),
+            int(ptc[1] + line_length*np.sin(heading)))
+        cv2.arrowedLine(vis, pt1=pt1, pt2=pt2, color=color,
+                        thickness=3, line_type=cv2.LINE_AA)
+
+    @staticmethod
+    def _draw_predicted_future(vis, pts, color):
+        # draw mean trajectory
+        for i in range(len(pts)-1):
+            pt1 = (int(pts[i][0]), int(pts[i][1]))
+            pt2 = (int(pts[i+1][0]), int(pts[i+1][1]))
+            cv2.line(vis, pt1=pt1, pt2=pt2, color=color,
+                    thickness=1, lineType=cv2.LINE_AA)
+
+    @staticmethod
+    def _draw_trust_region(vis, pts, r, color):
+        # draw trust regions
+        # vis: image
+        # pts: points
+        # r: radius
+        # color: color
+        for i in range(len(pts)):
+            ptc = (int(pts[i][0]), int(pts[i][1]))
+            rx, ry = int(r[i][0]), int(r[i][1])  # trust region
+            cv2.ellipse(vis, center=ptc, axes=(rx, ry), angle=0, startAngle=0, endAngle=360,
+                        color=color, thickness=1, lineType=cv2.LINE_AA)
+
+    @staticmethod
+    def _print_vehicle_info(vis, ptc, v, color):
+        # print latitude
+        pt = (ptc[0] + 15, ptc[1])
+        text = "%.6f" % v.x
+        cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                    fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+        # print longitude
+        pt = (ptc[0] + 15, ptc[1] + 20)
+        text = "%.6f" % v.y
+        cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                    fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+        # print id
+        if hasattr(v, 'traj_id') and v.traj_id is not None:
+            pt = (ptc[0] + 15, ptc[1] + 40)
+            text = f"id: {v.traj_id}"
+            cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                        fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+        if v.heading is not None:
+            pt= (ptc[0] + 15, ptc[1] + 60)
+            text = f"heading: {v.heading:.2f} deg"
+            cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                        fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+        # pt = (ptc[0] + 15, ptc[1] + 60)
+        # if hasattr(v, 'lane') and v.lane is not None:
+        #     text = f"lane: {v.lane}"
+        #     cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        #                 fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+        # # print category
+        # if v.category is not None:
+        #     pt = (ptc[0] + 15, ptc[1] + 40)
+        #     if label_list is not None:
+        #         text = "%s" % label_list.id2label[str(v.category)]
+        #     else:
+        #         text = "Category: %d" % v.category
+        #     cv2.putText(vis, text=text, org=pt, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        #                 fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
+
+    def _create_coord_mapper(self):
+
+        # wd_tl, wd_tr, wd_bl, wd_br = self.f.tl.copy(), self.f.tr.copy(), self.f.bl.copy(), self.f.br.copy()
+        px_tl, px_tr, px_bl, px_br = [0, 0], [
+            self.w-1, 0], [0, self.h-1], [self.w-1, self.h-1]
+
+        # normalize to local coordination
+        wd_tl = coord_normalization(
+            self.f.tl[0], self.f.tl[1], self.f.tl[0], self.f.tl[1])
+        wd_tr = coord_normalization(
+            self.f.tr[0], self.f.tr[1], self.f.tl[0], self.f.tl[1])
+        wd_bl = coord_normalization(
+            self.f.bl[0], self.f.bl[1], self.f.tl[0], self.f.tl[1])
+        wd_br = coord_normalization(
+            self.f.br[0], self.f.br[1], self.f.tl[0], self.f.tl[1])
+
+        wd_points = np.array([wd_tl, wd_tr, wd_bl, wd_br], np.float32)
+        px_points = np.array([px_tl, px_tr, px_bl, px_br], np.float32)
+
+        transform_wd2px = cv2.getPerspectiveTransform(
+            src=wd_points, dst=px_points)
+        transform_px2wd = cv2.getPerspectiveTransform(
+            src=px_points, dst=wd_points)
+
+        return transform_wd2px, transform_px2wd
+
+    def _world2pxl(self, pt_world, output_int=True):
+
+        # normalize to local coordination
+        pt_world = np.array(pt_world).reshape([-1, 2])
+        lat, lon = pt_world[:, 0], pt_world[:, 1]
+        lat_norm, lon_norm = coord_normalization(
+            lat, lon, self.f.tl[0], self.f.tl[1])
+        pt_world = np.array([lat_norm, lon_norm]).T
+
+        pt_world = np.array(pt_world).reshape([-1, 1, 2]).astype(np.float32)
+        pt_pixel = cv2.perspectiveTransform(pt_world, self.transform_wd2px)
+        pt_pixel = np.squeeze(pt_pixel)
+
+        if output_int:
+            pt_pixel = pt_pixel.astype(np.int32)
+
+        return pt_pixel
+
+    def _pxl2world(self, pt_pixel):
+
+        pt_pixel = np.array(pt_pixel).reshape([-1, 1, 2]).astype(np.float32)
+        pt_world = cv2.perspectiveTransform(
+            pt_pixel, self.transform_px2wd).astype(np.float64)
+        pt_world = np.squeeze(pt_world)
+        pt_world = pt_world.reshape([-1, 2])
+
+        # unnormalize to world coordination
+        lat_norm, lon_norm = pt_world[:, 0], pt_world[:, 1]
+        lat, lon = coord_unnormalization(
+            lat_norm, lon_norm, self.f.tl[0], self.f.tl[1])
+        pt_world = np.array([lat, lon]).T
+
+        return pt_world
+
+    def draw_points(self, frame, show_heading=False):
+        # TODO: draw vehicle as box when show_heading is True
+        if len(frame) == 0:
+            return self.basemap
+
+        vis = np.copy(self.basemap)
+
+        for i in range(0, len(frame)):
+            v = frame[i]
+
+            if v.x is None or v.y is None:
+                continue
+
+            color = self.color_table[hash(v.traj_id) % 10].tolist()
+
+            ptc = self._world2pxl([v.x, v.y])
+
+            # box unavailiable, draw a circle instead
+            self._draw_vehicle_as_point(vis, ptc, color)
+
+            if show_heading:
+                self._draw_vehicle_heading_as_arrow(vis, ptc, v.heading, color)
+
+            # print vehicle info beside box
+            self._print_vehicle_info(vis, ptc, v, (255, 255, 0),)
+
+        return vis
+
+    def draw_trajectory(self, frame, linewidth=2):
+
+        # update trajectory manager
+        for v in frame:
+            self.traj_manager.add_object(v, v.traj_id, frame.step, timestamp=frame.timestamp)
+
+        # update alpha (fade out)
+        self.traj_alpha *= 0.95
+
+        # draw trajectory
+        for v in self.traj_manager.last_frame:
+            pt_map = self._world2pxl([v.x, v.y])
+            v_ = v.prev
+            if v_ is None:
+                # print('prev is None')
+                continue
+            pt_map_prev = self._world2pxl([v_.x, v_.y])
+
+            pt1 = (int(pt_map[0]), int(pt_map[1]))
+            pt2 = (int(pt_map_prev[0]), int(pt_map_prev[1]))
+
+            color = self.color_table[hash(v.traj_id) % 10].tolist()
+
+            cv2.line(self.traj_layer, pt1=pt1, pt2=pt2,
+                        color=color, thickness=linewidth)
+            cv2.line(self.traj_alpha, pt1=pt1, pt2=pt2,
+                        color=(0.8, 0.8, 0.8), thickness=linewidth)
+
+        return self.traj_layer, self.traj_alpha
+
+    @staticmethod
+    def layer_blending(base_layer, traj_layer, traj_alpha):
+        base_layer = base_layer.astype(np.float32)/255.
+        traj_layer = traj_layer.astype(np.float32)/255.
+        vis = base_layer*(1-traj_alpha) + traj_layer*traj_alpha
+        return vis
+
+    def draw_polygon(self, base_layer, polygon, color, thickness=2):
+        pts = self._world2pxl(polygon)
+        layer = cv2.polylines(
+            base_layer, [pts], isClosed=True, color=color, thickness=thickness)
+        return layer
+
+    def render(self, frame,  with_traj=True, linewidth=2, show_heading=False):
+        base_layer = self.draw_points(
+            frame, show_heading=show_heading)
+        
+        if with_traj:
+            traj_layer, traj_alpha = self.draw_trajectory(
+                frame, linewidth)
+            map_vis = self.layer_blending(base_layer, traj_layer, traj_alpha)
+            map_vis = (map_vis*255.).astype(np.uint8)
+        else:
+            map_vis = base_layer
+
+        return map_vis
+    
+    def render_roaduser_points(self, list_of_points, show_heading=False):
+        ## this method is used to render a list of roaduser point, this is particularly useful when you have a list of points without knowing the id (so you can't generate frame)
+        
+        base_layer = self.draw_points(
+            list_of_points, show_heading=show_heading)
+        return base_layer
+

msight_base-0.1.2.dist-info/METADATA

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: msight_base
+Version: 0.1.2
+Summary: Object classes and base trajectory management for the MSight system
+Author-email: Rusheng Zhang <zrs1990@gmail.com>
+Requires-Python: >=3.6
+Requires-Dist: numpy
+Requires-Dist: opencv-python
+Requires-Dist: tqdm
+Requires-Dist: commonroad-io

msight_base-0.1.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+msight_base/__init__.py,sha256=pkeaGsGHCbbH4o0qkSjFb7eAKl-nMNiorNwYBNgYWEY,77
+msight_base/behavior.py,sha256=vEouWCXx1rjKnmRDgEFgDo8DPyoh_OZegNxXi42OH6s,732
+msight_base/detection.py,sha256=v6GG4UuLOyZcwdRz2dWUX-O7AWCMwvsT-_mhWchS7Nw,1643
+msight_base/map.py,sha256=YCM_oNYkmtmpeerlvbdreGsF0-mOOVHuzaO5TbF3-74,9251
+msight_base/road_user.py,sha256=X0US7aaY1a5-gUVYLGT3HV_QHDNaKEewMygpAlPdUEQ,6651
+msight_base/trajectory.py,sha256=eNq_dSrBVy4q7BVqqGnYVfyGc0zxLwNPMySn9DRgP-4,6894
+msight_base/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+msight_base/utils/cls.py,sha256=ShJ3loBdzEN9m4dqexU-hWQ60nV_c879siw5fbPAXOY,315
+msight_base/utils/data.py,sha256=ByvesNDSozRTCwQBAySrVLmavjgZH9UAwG2OfEnREx8,1194
+msight_base/visualizer/__init__.py,sha256=IEWq4Rah8uNuhp9ay1lj3t8n40tY-ZDc3g3DKxZWv0I,36
+msight_base/visualizer/utils.py,sha256=kd5OwLWELhGK5N-cCv0nDS-mHApPMcQ9erQMmAQzM-U,4297
+msight_base/visualizer/visualizer.py,sha256=lBFIyfETqFaXCMK2NdLkg-4MGkVmrGSwtf6XFWLgDwM,11324
+msight_base-0.1.2.dist-info/METADATA,sha256=GMAQnxBCpksvEUbSqLB6Ukcex1Dye9KYD3lNXm7MbLs,312
+msight_base-0.1.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+msight_base-0.1.2.dist-info/top_level.txt,sha256=6I_xkTYR6LZXAjnAiTFNVItZhpdUpTRMj_TjvOs2Wig,12
+msight_base-0.1.2.dist-info/RECORD,,

msight_base-0.1.2.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

msight_base-0.1.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+msight_base