diff --git a/.gitignore b/.gitignore index 2af531ab1..f02ef2e3c 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ *.log *.csv *.ipynb +*.dot dependencies/ pylot/prediction/prediction/ diff --git a/data_gatherer.py b/data_gatherer.py index 0291f360f..46221b111 100644 --- a/data_gatherer.py +++ b/data_gatherer.py @@ -162,7 +162,8 @@ def main(argv): traffic_light_depth_camera_stream, traffic_light_segmented_camera_stream, pose_stream) - pylot.operator_creator.add_bounding_box_logging(traffic_lights_stream) + pylot.operator_creator.add_bounding_box_logging( + traffic_lights_stream, 'tl-bboxes') if FLAGS.log_left_right_cameras: (left_camera_stream, right_camera_stream, _, @@ -186,7 +187,8 @@ def main(argv): depth_camera_stream, center_camera_stream, segmented_stream, pose_stream, ground_obstacles_stream, ground_speed_limit_signs_stream, ground_stop_signs_stream) - pylot.operator_creator.add_bounding_box_logging(obstacles_stream) + pylot.operator_creator.add_bounding_box_logging( + obstacles_stream, 'bboxes') if FLAGS.log_multiple_object_tracker: pylot.operator_creator.add_multiple_object_tracker_logging( @@ -300,8 +302,6 @@ def main(argv): stream.send(erdos.WatermarkMessage(erdos.Timestamp(is_top=True))) try: - if pylot.flags.must_visualize(): - pylot.utils.run_visualizer_control_loop(control_display_stream) node_handle.wait() except KeyboardInterrupt: node_handle.shutdown() diff --git a/docker/Dockerfile b/docker/Dockerfile index 80991608b..f7c092cab 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -72,3 +72,35 @@ RUN echo "if [ -f ~/.bashrc ]; then . ~/.bashrc ; fi" >> ~/.bash_profile # Set up ssh access to the container. RUN cd ~/ && ssh-keygen -q -t rsa -N '' -f ~/.ssh/id_rsa <<&1 >/dev/null RUN sudo sed -i 's/#X11UseLocalhost yes/X11UseLocalhost no/g' /etc/ssh/sshd_config + +# WIP (for visualization in Foxglove) + +# Set up for Foxglove +RUN sudo apt-get -y install wget + +# Install Foxglove Studio +RUN cd /tmp +RUN wget https://github.com/foxglove/studio/releases/download/v0.12.0/foxglove-studio-0.12.0-linux-amd64.deb +RUN sudo apt install -y ./foxglove-studio-*.deb +RUN rm ./foxglove-studio-*.deb + +# Set up for ROS +RUN sudo apt -y update +RUN sudo apt install -y lsb-release +RUN sudo apt install -y curl + +# Download ROS +# From, http://wiki.ros.org/melodic/Installation/Ubuntu: +# Setup your computer to accept software from packages.ros.org: +RUN sudo sh -c 'echo "deb http://packages.ros.org/ros/ubuntu $(lsb_release -sc) main" > /etc/apt/sources.list.d/ros-latest.list' +# Set up your keys +# RUN sudo apt-key adv --keyserver 'hkp://keyserver.ubuntu.com:80' --recv-key C1CF6E31E6BADE8868B172B4F42ED6FBAB17C654 +RUN curl -s https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc | sudo apt-key add - +# Make sure your Debian package index is up-to-date +RUN sudo apt -y update +# Desktop-Full Install +RUN sudo apt -y install ros-melodic-desktop-full +# Set up ROS environment variables +RUN echo "source /opt/ros/melodic/setup.bash" >> ~/.bashrc +RUN source ~/.bashrc + diff --git a/install.sh b/install.sh index fca7f4cc1..632c0675a 100755 --- a/install.sh +++ b/install.sh @@ -78,6 +78,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 @@ -140,6 +145,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/ diff --git a/pylot.py b/pylot.py index 572712e8f..cbd52793c 100644 --- a/pylot.py +++ b/pylot.py @@ -218,7 +218,7 @@ def driver(): prediction_stream, waypoints_stream, control_stream) streams_to_send_top_on += ingest_streams - node_handle = erdos.run_async() + node_handle = erdos.run_async('pylot.dot') for stream in streams_to_send_top_on: stream.send(erdos.WatermarkMessage(erdos.Timestamp(is_top=True))) @@ -260,8 +260,6 @@ def main(args): client.start_recorder(FLAGS.simulation_recording_file) node_handle, control_display_stream = driver() signal.signal(signal.SIGINT, shutdown) - if pylot.flags.must_visualize(): - pylot.utils.run_visualizer_control_loop(control_display_stream) node_handle.wait() except KeyboardInterrupt: shutdown_pylot(node_handle, client, world) diff --git a/pylot/component_creator.py b/pylot/component_creator.py index 8df37629c..d4bb474c9 100644 --- a/pylot/component_creator.py +++ b/pylot/component_creator.py @@ -354,6 +354,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 = \ diff --git a/pylot/debug/ros_camera_publisher.py b/pylot/debug/ros_camera_publisher.py new file mode 100644 index 000000000..83c90ada9 --- /dev/null +++ b/pylot/debug/ros_camera_publisher.py @@ -0,0 +1,34 @@ +import sys +import rospy +import numpy as np +from sensor_msgs.msg import Image + + +class ROSCameraPublisher: + """Class that stores a ROS publisher node that publishes ROS Image messages + + Args: + topic: the name of the topic published to + + Attributes: + image_pub: ROS publisher node + """ + def __init__(self, topic: str): + self.image_pub = rospy.Publisher(topic, Image, queue_size=10) + + def publish(self, img_arr): + """Publishes a sensor_msgs/Image message (constructed from input) + + Args: + img_arr: A numpy array storing a frame (e.g. camera, depth, etc.) + """ + img_msg = Image(encoding='rgb8') + if type(img_arr[0][0][0]) != np.int8: + img_arr = img_arr.astype(np.int8) + img_msg.height, img_msg.width, channels = img_arr.shape + img_msg.data = img_arr.tobytes() + img_msg.step = img_msg.width * img_msg.height + img_msg.is_bigendian = (img_arr.dtype.byteorder == '>' + or img_arr.dtype.byteorder == '=' + and sys.byteorder == 'big') + self.image_pub.publish(img_msg) diff --git a/pylot/debug/ros_lidar_publisher.py b/pylot/debug/ros_lidar_publisher.py new file mode 100644 index 000000000..37449887a --- /dev/null +++ b/pylot/debug/ros_lidar_publisher.py @@ -0,0 +1,45 @@ +import sys +import rospy +import numpy as np +from sensor_msgs.msg import PointCloud2, PointField + + +class ROSLIDARPublisher: + """Class that stores a ROS publisher node that publishes ROS Point Cloud messages + + Args: + topic: the name of the topic published to + + Attributes: + point_cloud_pub: ROS publisher node + """ + def __init__(self, topic: str): + self.point_cloud_pub = rospy.Publisher(topic, + PointCloud2, + queue_size=10) + + def publish(self, points): + """Publishes a sensor_msgs/PointCloud2 message (constructed from input) + + Args: + points: A numpy array storing a point cloud (see pylot.pylot.perception.point_cloud) + """ + + points = points.astype(np.float32) + points_byte_array = points.tobytes() + row_step = len(points_byte_array) + point_step = len(points[0].tobytes()) + fields = [ + PointField('x', 0, PointField.FLOAT32, 1), + PointField('y', 4, PointField.FLOAT32, 1), + PointField('z', 8, PointField.FLOAT32, 1) + ] + point_cloud_msg = PointCloud2(height=1, + width=len(points), + is_dense=True, + is_bigendian=False, + fields=fields, + point_step=point_step, + row_step=row_step, + data=points_byte_array) + self.point_cloud_pub.publish(point_cloud_msg) diff --git a/pylot/debug/visualizer_operator.py b/pylot/debug/visualizer_operator.py index ce069bf3a..0cea2ada9 100644 --- a/pylot/debug/visualizer_operator.py +++ b/pylot/debug/visualizer_operator.py @@ -8,24 +8,24 @@ import erdos import numpy as np - -import pygame -from pygame.locals import K_n +import rospy +import subprocess import pylot.utils from pylot.drivers.sensor_setup import RGBCameraSetup from pylot.perception.camera_frame import CameraFrame from pylot.planning.world import World +from pylot.debug.ros_camera_publisher import ROSCameraPublisher +from pylot.debug.ros_lidar_publisher import ROSLIDARPublisher DEFAULT_VIS_TIME = 30000.0 - class VisualizerOperator(erdos.Operator): """ The `VisualizerOperator` allows developers to see the current state - of the entire pipeline by visualizing it on a pygame instance. + of the entire pipeline by visualizing it in Foxglove Studio. This receives input data from almost the entire pipeline and renders the - results of the operator currently chosen by the developer on the screen. + results of the operator(s) in Foxglove Studio. """ def __init__(self, pose_stream, rgb_camera_stream, tl_camera_stream, prediction_camera_stream, depth_camera_stream, @@ -33,7 +33,7 @@ def __init__(self, pose_stream, rgb_camera_stream, tl_camera_stream, obstacles_stream, traffic_lights_stream, tracked_obstacles_stream, lane_detection_stream, prediction_stream, waypoints_stream, control_stream, - display_control_stream, pygame_display, flags): + control_display_stream, flags): visualize_streams = [] self._pose_msgs = deque() pose_stream.add_callback( @@ -129,52 +129,38 @@ def __init__(self, pose_stream, rgb_camera_stream, tl_camera_stream, erdos.add_watermark_callback(visualize_streams, [], self.on_watermark) # Add a callback on a control stream to figure out what to display. - display_control_stream.add_callback(self.change_display) self._logger = erdos.utils.setup_logging(self.config.name, self.config.log_file_name) - self.display = pygame_display - - # Set the font. - fonts = [x for x in pygame.font.get_fonts() if 'mono' in x] - default_font = 'ubuntumono' - mono = default_font if default_font in fonts else fonts[0] - mono = pygame.font.match_font(mono) - self.font = pygame.font.Font(mono, 14) - - # Array of keys to figure out which message to display. - self.current_display = 0 - self.display_array = [] - self.window_titles = [] + + # Pylot-ROS Integration for Foxglove visualization + self.roscore = subprocess.Popen('roscore') + import time + time.sleep(1) + rospy.init_node("visualizer", anonymous=True, disable_signals=True) + self.pub = {} # dict of publishers + + # Creating ROS publishers for streams to be visualized if flags.visualize_rgb_camera: - self.display_array.append("RGB") - self.window_titles.append("RGB Camera") + self.pub["RGB"] = ROSCameraPublisher("/camera/rgb") if flags.visualize_detected_obstacles: - self.display_array.append("Obstacle") - self.window_titles.append("Detected obstacles") + self.pub["Obstacle"] = ROSCameraPublisher("/camera/obstacle") if flags.visualize_tracked_obstacles: - self.display_array.append("TrackedObstacle") - self.window_titles.append("Obstacle tracking") + self.pub["TrackedObstacle"] = ROSCameraPublisher( + "/camera/tracked_obstacle") if flags.visualize_detected_traffic_lights: - self.display_array.append("TLCamera") - self.window_titles.append("Detected traffic lights") + self.pub["TLCamera"] = ROSCameraPublisher("/camera/tl_camera") if flags.visualize_waypoints: - self.display_array.append("Waypoint") - self.window_titles.append("Planning") - if flags.visualize_prediction: - self.display_array.append("PredictionCamera") - self.window_titles.append("Prediction") + self.pub["Waypoint"] = ROSCameraPublisher("/camera/waypoint") if flags.visualize_lidar: - self.display_array.append("PointCloud") - self.window_titles.append("LiDAR") + self.pub["PointCloud"] = ROSLIDARPublisher("/lidar/point_cloud") if flags.visualize_detected_lanes: - self.display_array.append("Lanes") - self.window_titles.append("Detected lanes") + self.pub["Lanes"] = ROSCameraPublisher("/camera/lanes") if flags.visualize_depth_camera: - self.display_array.append("Depth") - self.window_titles.append("Depth Camera") + self.pub["Depth"] = ROSCameraPublisher("/camera/depth") + self.pub["DepthPointCloud"] = ROSLIDARPublisher("/lidar/depth") if flags.visualize_segmentation: - self.display_array.append("Segmentation") - self.window_titles.append("Segmentation") + self.pub["Segmentation"] = ROSCameraPublisher( + "/camera/segmentation") if flags.visualize_world: self._planning_world = World(flags, self._logger) top_down_transform = pylot.utils.get_top_down_transform( @@ -184,12 +170,10 @@ def __init__(self, pose_stream, rgb_camera_stream, tl_camera_stream, self._bird_eye_camera_setup = RGBCameraSetup( 'bird_eye_camera', flags.camera_image_width, flags.camera_image_height, top_down_transform, 90) - self.display_array.append("PlanningWorld") - self.window_titles.append("Planning world") + self.pub["PlanningWorld"] = ROSCameraPublisher( + "/camera/planning_world") else: self._planning_world = None - assert len(self.display_array) == len(self.window_titles), \ - "The display and titles differ." # Save the flags. self._flags = flags @@ -205,19 +189,13 @@ def connect(pose_stream, rgb_camera_stream, tl_camera_stream, def destroy(self): self._logger.warn('destroying {}'.format(self.config.name)) + self.roscore.terminate() def save(self, msg, msg_type, queue): self._logger.debug("@{}: Received {} message.".format( msg.timestamp, msg_type)) queue.append(msg) - def change_display(self, display_message): - if display_message.data == K_n: - self.current_display = (self.current_display + 1) % len( - self.display_array) - self._logger.debug("@{}: Visualizer changed to {}".format( - display_message.timestamp, self.current_display)) - def get_message(self, queue, timestamp, name): msg = None if queue: @@ -258,21 +236,7 @@ def render_text(self, pose, control, timestamp): "Reverse : {:.2f}".format(control.reverse), ] - # Display the information box. - info_surface = pygame.Surface( - (220, self._flags.camera_image_height // 3)) - info_surface.set_alpha(100) - self.display.blit(info_surface, (0, 0)) - - # Render the text. - v_offset = 10 - for line in info_text: - if v_offset + 18 > self._flags.camera_image_height: - break - surface = self.font.render(line, True, (255, 255, 255)) - self.display.blit(surface, (8, v_offset)) - v_offset += 18 - pygame.display.flip() + # TODO: Render text (using ROS publisher) def on_watermark(self, timestamp): self._logger.debug("@{}: received watermark.".format(timestamp)) @@ -315,27 +279,27 @@ def on_watermark(self, timestamp): if self._flags.visualize_imu: self._visualize_imu(imu_msg) - sensor_to_display = self.display_array[self.current_display] - if sensor_to_display == "RGB" and bgr_msg: - bgr_msg.frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "Obstacle" and bgr_msg and obstacle_msg: + if self._flags.visualize_rgb_camera and bgr_msg: + image_np = bgr_msg.frame.as_rgb_numpy_array() + self.pub["RGB"].publish(image_np) + if self._flags.visualize_detected_obstacles and bgr_msg and obstacle_msg: bgr_msg.frame.annotate_with_bounding_boxes(timestamp, obstacle_msg.obstacles, ego_transform) - bgr_msg.frame.visualize(self.display, timestamp=timestamp) - elif (sensor_to_display == "TLCamera" and tl_camera_msg - and traffic_light_msg): + image_np = bgr_msg.frame.as_rgb_numpy_array() + self.pub["Obstacle"].publish(image_np) + if self._flags.visualize_detected_traffic_lights and tl_camera_msg and traffic_light_msg: tl_camera_msg.frame.annotate_with_bounding_boxes( timestamp, traffic_light_msg.obstacles) - tl_camera_msg.frame.visualize(self.display, timestamp=timestamp) - elif (sensor_to_display == "TrackedObstacle" and bgr_msg - and tracked_obstacle_msg): + image_np = tl_camera_msg.frame.as_rgb_numpy_array() + self.pub["TLCamera"].publish(image_np) + if self._flags.visualize_tracked_obstacles and bgr_msg and tracked_obstacle_msg: bgr_msg.frame.annotate_with_bounding_boxes( timestamp, tracked_obstacle_msg.obstacle_trajectories, ego_transform) - bgr_msg.frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "Waypoint" and (bgr_msg and pose_msg - and waypoint_msg): + image_np = bgr_msg.frame.as_rgb_numpy_array() + self.pub["TrackedObstacle"].publish(image_np) + if self._flags.visualize_waypoints and bgr_msg and pose_msg and waypoint_msg: bgr_frame = bgr_msg.frame if self._flags.draw_waypoints_on_camera_frames: bgr_frame.camera_setup.set_transform( @@ -343,26 +307,39 @@ def on_watermark(self, timestamp): waypoint_msg.waypoints.draw_on_frame(bgr_frame) if self._flags.draw_waypoints_on_world: waypoint_msg.waypoints.draw_on_world(self._world) - bgr_frame.visualize(self.display, timestamp=timestamp) - elif (sensor_to_display == "PredictionCamera" and prediction_camera_msg - and prediction_msg): + image_np = bgr_frame.as_rgb_numpy_array() + self.pub["Waypoint"].publish(image_np) + if prediction_camera_msg and prediction_msg: frame = prediction_camera_msg.frame frame.transform_to_cityscapes() for obstacle_prediction in prediction_msg.predictions: obstacle_prediction.draw_trajectory_on_frame(frame) - frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "PointCloud" and point_cloud_msg: - point_cloud_msg.point_cloud.visualize(self.display, - timestamp=timestamp) - elif (sensor_to_display == "Lanes" and bgr_msg and lane_detection_msg): + if self._flags.visualize_lidar and point_cloud_msg: + points = point_cloud_msg.point_cloud.points + # need to switch and reverse axes (from Velodyne coordinate space) + # for correct visualization in foxglove + points[:, [0, 2]] = points[:, [2, 0]] + points[:, [1, 2]] = points[:, [2, 1]] + points[:, 0] = -points[:, 0] + points[:, 2] = -points[:, 2] + self.pub["PointCloud"].publish(points) + if self._flags.visualize_detected_lanes and bgr_msg and lane_detection_msg: for lane in lane_detection_msg.data: lane.draw_on_frame(bgr_msg.frame) - bgr_msg.frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "Depth" and depth_msg: - depth_msg.frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "Segmentation" and segmentation_msg: - segmentation_msg.frame.visualize(self.display, timestamp=timestamp) - elif sensor_to_display == "PlanningWorld": + image_np = bgr_msg.frame.as_rgb_numpy_array() + self.pub["Lanes"].publish(image_np) + if self._flags.visualize_depth_camera and depth_msg: + image_np = depth_msg.frame.original_frame + image_np = image_np[:, :, ::-1] + self.pub["Depth"].publish(image_np) + depth_msg.frame.resize(854, 480) + points = depth_msg.frame.as_point_cloud() + points[:, 0] = -points[:, 0] + self.pub["DepthPointCloud"].publish(points) + if self._flags.visualize_segmentation and segmentation_msg: + self.pub["Segmentation"].publish( + segmentation_msg.frame.as_cityscapes_palette()) + if self._flags.visualize_world: if prediction_camera_msg is None: # Top-down prediction is not available. Show planning # world on a black image. @@ -386,9 +363,11 @@ def on_watermark(self, timestamp): lanes=lanes) self._planning_world.update_waypoints(None, waypoint_msg.waypoints) self._planning_world.draw_on_frame(frame) - frame.visualize(self.display, timestamp=timestamp) + image_np = frame.as_rgb_numpy_array() + self.pub["PlanningWorld"].publish(image_np) - self.render_text(pose_msg.data, control_msg, timestamp) + # TODO: render text in foxglove + # self.render_text(pose_msg.data, control_msg, timestamp) def run(self): # Run method is invoked after all operators finished initializing. diff --git a/pylot/flags.py b/pylot/flags.py index bbf482345..7b88c7f71 100644 --- a/pylot/flags.py +++ b/pylot/flags.py @@ -40,9 +40,10 @@ 'True to enable obstacle tracking operator') 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'], - 'Sets which obstacle tracker to use') +flags.DEFINE_enum( + 'tracker_type', 'sort', + ['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, 'True to enable perfect lane detection') diff --git a/pylot/localization/localization_operator.py b/pylot/localization/localization_operator.py index 66c7f5d2b..9974aaa79 100644 --- a/pylot/localization/localization_operator.py +++ b/pylot/localization/localization_operator.py @@ -182,10 +182,8 @@ def on_watermark(self, timestamp: Timestamp): if pose_msg: self._last_pose_estimate = pose_msg.data self._last_timestamp = timestamp.coordinates[0] - if (self._flags.execution_mode == 'challenge-map' - or self._flags.execution_mode == 'challenge-sensors'): - self._pose_stream.send(pose_msg) - self._pose_stream.send(erdos.WatermarkMessage(timestamp)) + self._pose_stream.send(pose_msg) + self._pose_stream.send(erdos.WatermarkMessage(timestamp)) else: raise NotImplementedError( "Need pose message to initialize the estimates.") diff --git a/pylot/loggers/bounding_box_logger_operator.py b/pylot/loggers/bounding_box_logger_operator.py index 1618f9bd2..e6225dcc9 100644 --- a/pylot/loggers/bounding_box_logger_operator.py +++ b/pylot/loggers/bounding_box_logger_operator.py @@ -21,13 +21,15 @@ class BoundingBoxLoggerOperator(erdos.Operator): _data_path (:obj:`str`): Directory to which to log files. """ def __init__(self, obstacles_stream: erdos.ReadStream, - finished_indicator_stream: erdos.WriteStream, flags): + finished_indicator_stream: erdos.WriteStream, flags, + file_base_name): obstacles_stream.add_callback(self.on_obstacles_msg) self._logger = erdos.utils.setup_logging(self.config.name, self.config.log_file_name) self._flags = flags + self._file_base_name = file_base_name self._msg_cnt = 0 - self._data_path = os.path.join(self._flags.data_path, 'bboxes') + self._data_path = os.path.join(self._flags.data_path, file_base_name) os.makedirs(self._data_path, exist_ok=True) @staticmethod @@ -57,7 +59,8 @@ def on_obstacles_msg(self, msg: erdos.Message): assert len(msg.timestamp.coordinates) == 1 timestamp = msg.timestamp.coordinates[0] # Write the bounding boxes. - file_name = os.path.join(self._data_path, - 'bboxes-{}.json'.format(timestamp)) + file_name = os.path.join( + self._data_path, '{}-{}.json'.format(self._file_base_name, + timestamp)) with open(file_name, 'w') as outfile: json.dump(bboxes, outfile) diff --git a/pylot/operator_creator.py b/pylot/operator_creator.py index c011c28a9..b4d235874 100644 --- a/pylot/operator_creator.py +++ b/pylot/operator_creator.py @@ -269,6 +269,19 @@ 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, @@ -662,6 +675,7 @@ def add_planning_pose_synchronizer(waypoint_stream, pose_stream, def add_bounding_box_logging(obstacles_stream, + file_base_name, name='bounding_box_logger_operator'): from pylot.loggers.bounding_box_logger_operator import \ BoundingBoxLoggerOperator @@ -671,7 +685,7 @@ def add_bounding_box_logging(obstacles_stream, profile_file_name=FLAGS.profile_file_name) [finished_indicator_stream] = erdos.connect(BoundingBoxLoggerOperator, op_config, [obstacles_stream], - FLAGS) + FLAGS, file_base_name) return finished_indicator_stream @@ -823,12 +837,6 @@ def add_visualizer(pose_stream=None, control_stream=None, name='visualizer_operator'): from pylot.debug.visualizer_operator import VisualizerOperator - import pygame - pygame.init() - pygame_display = pygame.display.set_mode( - (FLAGS.camera_image_width, FLAGS.camera_image_height), - pygame.HWSURFACE | pygame.DOUBLEBUF) - pygame.display.set_caption("Pylot") streams_to_send_top_on = [] if pose_stream is None: pose_stream = erdos.IngestStream() @@ -897,7 +905,7 @@ def add_visualizer(pose_stream=None, obstacles_stream, traffic_lights_stream, tracked_obstacles_stream, lane_detection_stream, prediction_stream, waypoints_stream, control_stream, control_display_stream - ], pygame_display, FLAGS) + ], FLAGS) return control_display_stream, streams_to_send_top_on diff --git a/pylot/perception/camera_frame.py b/pylot/perception/camera_frame.py index 3e761f7e1..9bfee8602 100644 --- a/pylot/perception/camera_frame.py +++ b/pylot/perception/camera_frame.py @@ -137,19 +137,6 @@ def resize(self, width: int, height: int): dsize=(width, height), interpolation=cv2.INTER_NEAREST) - def visualize(self, pygame_display, timestamp=None): - """Visualizes the frame on a pygame display.""" - import pygame - if timestamp is not None: - pylot.utils.add_timestamp(self.frame, timestamp) - if self.encoding != 'RGB': - image_np = self.as_rgb_numpy_array() - else: - image_np = self.frame - image_np = np.transpose(image_np, (1, 0, 2)) - pygame.surfarray.blit_array(pygame_display, image_np) - pygame.display.flip() - def save(self, timestamp: int, data_path: str, file_base: str): """Saves the camera frame to a file. diff --git a/pylot/perception/depth_frame.py b/pylot/perception/depth_frame.py index bf49987da..084ba2cae 100644 --- a/pylot/perception/depth_frame.py +++ b/pylot/perception/depth_frame.py @@ -128,16 +128,6 @@ def resize(self, width: int, height: int): dsize=(width, height), interpolation=cv2.INTER_NEAREST) - def visualize(self, pygame_display, timestamp=None): - """Visualizes the frame on a pygame display.""" - if self.original_frame is not None: - import pygame - image_np = self.original_frame - image_np = image_np[:, :, ::-1] - image_np = np.transpose(image_np, (1, 0, 2)) - pygame.surfarray.blit_array(pygame_display, image_np) - pygame.display.flip() - def save(self, timestamp: int, data_path: str, file_base: str): """Saves the depth frame to a file. diff --git a/pylot/perception/detection/lane.py b/pylot/perception/detection/lane.py index 3cc3e2303..70069d83d 100644 --- a/pylot/perception/detection/lane.py +++ b/pylot/perception/detection/lane.py @@ -5,7 +5,6 @@ import numpy as np from pylot.utils import Location, Rotation, Transform, Vector3D - from shapely.geometry import Point from shapely.geometry.polygon import Polygon diff --git a/pylot/perception/flags.py b/pylot/perception/flags.py index 4ed856922..a72b94a8b 100644 --- a/pylot/perception/flags.py +++ b/pylot/perception/flags.py @@ -84,6 +84,15 @@ ['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', 'Path to the model') +flags.DEFINE_string( + 'qd_track_config_path', + 'dependencies/qdtrack/configs/qdtrack-frcnn_r50_fpn_12e_bdd100k.py', + 'Path to the model') + # Lane detection flags. flags.DEFINE_float('lane_detection_gpu_memory_fraction', 0.3, 'GPU memory fraction allocated to Lanenet') diff --git a/pylot/perception/point_cloud.py b/pylot/perception/point_cloud.py index 35c40997e..3c4263227 100644 --- a/pylot/perception/point_cloud.py +++ b/pylot/perception/point_cloud.py @@ -184,24 +184,6 @@ def save(self, timestamp: int, data_path: str, file_base: str): pcd.points = o3d.Vector3dVector(self.points) o3d.write_point_cloud(file_name, pcd) - def visualize(self, pygame_display, timestamp=None): - """Visualizes the point cloud on a pygame display.""" - import pygame - (width, height) = pygame_display.get_size() - # Transform point cloud to top down view. - lidar_data = np.array(self.global_points[:, :2]) - lidar_data *= (min(width, height) / - (2.0 * self._lidar_setup.get_range_in_meters())) - lidar_data += (0.5 * width, 0.5 * height) - lidar_data = np.fabs(lidar_data) - lidar_data = lidar_data.astype(np.int32) - lidar_data = np.reshape(lidar_data, (-1, 2)) - lidar_img_size = (width, height, 3) - lidar_img = np.zeros((lidar_img_size), dtype=np.uint8) - lidar_img[tuple(lidar_data.T)] = (255, 255, 255) - pygame.surfarray.blit_array(pygame_display, lidar_img) - pygame.display.flip() - def __repr__(self): return 'PointCloud(lidar setup: {}, points: {})'.format( self._lidar_setup, self.points) diff --git a/pylot/perception/segmentation/segmented_frame.py b/pylot/perception/segmentation/segmented_frame.py index 2d793fe29..f87cf71d9 100644 --- a/pylot/perception/segmentation/segmented_frame.py +++ b/pylot/perception/segmentation/segmented_frame.py @@ -274,15 +274,6 @@ def save(self, timestamp, data_path, file_base): img = Image.fromarray(self.as_cityscapes_palette()) img.save(file_name) - def visualize(self, pygame_display, timestamp=None): - import pygame - cityscapes_frame = self.as_cityscapes_palette() - if timestamp is not None: - pylot.utils.add_timestamp(cityscapes_frame, timestamp) - image_np = np.transpose(cityscapes_frame, (1, 0, 2)) - pygame.surfarray.blit_array(pygame_display, image_np) - pygame.display.flip() - def draw_box(self, start_point, end_point, color, thickness=3): """Draw a colored box defined by start_point, end_point.""" start = (int(start_point.x), int(start_point.y)) diff --git a/pylot/perception/tracking/qd_track_operator.py b/pylot/perception/tracking/qd_track_operator.py new file mode 100644 index 000000000..2439c78d0 --- /dev/null +++ b/pylot/perception/tracking/qd_track_operator.py @@ -0,0 +1,74 @@ +import time + +import erdos + +from pylot.perception.detection.obstacle import Obstacle +from pylot.perception.detection.utils import BoundingBox2D, \ + OBSTACLE_LABELS +from pylot.perception.messages import ObstaclesMessage + + +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() + obstacles = [] + for k, v in track_result.items(): + track_id = k + bbox = v['bbox'][None, :] + score = bbox[4] + label_id = v['label'] + 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]) + 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)) diff --git a/pylot/simulation/carla_operator.py b/pylot/simulation/carla_operator.py index 0bc4bf31e..da59b7ea5 100644 --- a/pylot/simulation/carla_operator.py +++ b/pylot/simulation/carla_operator.py @@ -110,7 +110,9 @@ def __init__( # Spawn ego vehicle, people and vehicles. (self._ego_vehicle, self._vehicle_ids, self._people) = pylot.simulation.utils.spawn_actors( - self._client, self._world, self._simulator_version, + self._client, self._world, + self._flags.carla_traffic_manager_port, + self._simulator_version, self._flags.simulator_spawn_point_index, self._flags.control == 'simulator_auto_pilot', self._flags.simulator_num_people, diff --git a/pylot/simulation/perfect_lane_detector_operator.py b/pylot/simulation/perfect_lane_detector_operator.py index e4b12103b..81a9b5ba0 100644 --- a/pylot/simulation/perfect_lane_detector_operator.py +++ b/pylot/simulation/perfect_lane_detector_operator.py @@ -36,8 +36,6 @@ def __init__(self, pose_stream: ReadStream, open_drive_stream: ReadStream, self._flags = flags self._logger = erdos.utils.setup_logging(self.config.name, self.config.log_file_name) - self._logger_lane = erdos.utils.setup_logging( - self.config.name + "_lane", self.config.log_file_name + "_lane") self._bgr_msgs = deque() self._pose_msgs = deque() self._frame_cnt = 0 @@ -132,7 +130,7 @@ def on_position_update(self, timestamp: Timestamp, binary_img = Image.fromarray(binary_frame) instance_img.save(instance_file_name) binary_img.save(binary_file_name) - self._logger_lane.debug( + self._logger.debug( '@{}: Created binary lane and lane images in {}'. format(pose_msg.timestamp, self._flags.data_path)) else: diff --git a/pylot/simulation/utils.py b/pylot/simulation/utils.py index 653d00c84..214aab1a1 100644 --- a/pylot/simulation/utils.py +++ b/pylot/simulation/utils.py @@ -166,11 +166,13 @@ def reset_world(world): actor.destroy() -def spawn_actors(client, world, simulator_version: str, - ego_spawn_point_index: int, auto_pilot: bool, num_people: int, - num_vehicles: int, logger): - vehicle_ids = spawn_vehicles(client, world, num_vehicles, logger) - ego_vehicle = spawn_ego_vehicle(world, ego_spawn_point_index, auto_pilot) +def spawn_actors(client, world, traffic_manager_port: int, + simulator_version: str, ego_spawn_point_index: int, + auto_pilot: bool, num_people: int, num_vehicles: int, logger): + vehicle_ids = spawn_vehicles(client, world, traffic_manager_port, + num_vehicles, logger) + ego_vehicle = spawn_ego_vehicle(world, traffic_manager_port, + ego_spawn_point_index, auto_pilot) people = [] if check_simulator_version(simulator_version, @@ -189,6 +191,7 @@ def spawn_actors(client, world, simulator_version: str, def spawn_ego_vehicle(world, + traffic_manager_port: int, spawn_point_index: int, auto_pilot: bool, blueprint: str = 'vehicle.lincoln.mkz2017'): @@ -207,7 +210,7 @@ def spawn_ego_vehicle(world, ego_vehicle = world.try_spawn_actor(v_blueprint, start_pose) if auto_pilot: - ego_vehicle.set_autopilot(True) + ego_vehicle.set_autopilot(True, traffic_manager_port) return ego_vehicle @@ -271,7 +274,8 @@ def spawn_people(client, world, num_people: int, logger): return (ped_ids, ped_control_ids) -def spawn_vehicles(client, world, num_vehicles: int, logger): +def spawn_vehicles(client, world, traffic_manager_port: int, num_vehicles: int, + logger): """ Spawns vehicles at random locations inside the world. Args: @@ -309,7 +313,8 @@ def spawn_vehicles(client, world, num_vehicles: int, logger): batch.append( command.SpawnActor(blueprint, transform).then( - command.SetAutopilot(command.FutureActor, True))) + command.SetAutopilot(command.FutureActor, True, + traffic_manager_port))) # Apply the batch and retrieve the identifiers. vehicle_ids = [] diff --git a/pylot/utils.py b/pylot/utils.py index d3774e130..cffe6c436 100644 --- a/pylot/utils.py +++ b/pylot/utils.py @@ -1049,33 +1049,6 @@ def set_tf_loglevel(level): logging.getLogger('tensorflow').setLevel(level) -def run_visualizer_control_loop(control_display_stream): - """Runs a pygame loop that waits for user commands. - - The user commands are send on the control_display_stream - to control the pygame visualization window. - """ - import erdos - import pygame - clock = pygame.time.Clock() - from pygame.locals import K_n - while True: - clock.tick_busy_loop(60) - events = pygame.event.get() - for event in events: - if event.type == pygame.KEYUP: - if event.key == K_n: - control_display_stream.send( - erdos.Message(erdos.Timestamp(coordinates=[0]), - event.key)) - elif event.type == pygame.QUIT: - raise KeyboardInterrupt - elif event.type == pygame.KEYDOWN: - if (event.key == pygame.K_c - and pygame.key.get_mods() & pygame.KMOD_CTRL): - raise KeyboardInterrupt - - def verify_keys_in_dict(required_keys, arg_dict): assert set(required_keys).issubset(set(arg_dict.keys())), \ "one or more of {} not found in {}".format(required_keys, arg_dict) diff --git a/requirements.txt b/requirements.txt index 5db6b07f9..c300a9187 100644 --- a/requirements.txt +++ b/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