def __init__(self,
obstacles_stream,
camera_stream,
obstacle_tracking_stream,
name,
tracker_type,
flags,
log_file_name=None,
csv_file_name=None):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
self._name = name
self._flags = flags
self._logger = erdos.utils.setup_logging(name, log_file_name)
self._csv_logger = erdos.utils.setup_csv_logging(
name + '-csv', csv_file_name)
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import\
MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
# Labels the obstacle trackers should track.
self._tracked_labels = {
'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck'
}
# True when the tracker is ready to update bboxes.
self._ready_to_update = False
self._ready_to_update_timestamp = None
self._to_process = deque()
def __init__(self, obstacles_stream, camera_stream,
time_to_decision_stream, obstacle_tracking_stream,
tracker_type, flags):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
time_to_decision_stream.add_callback(self.on_time_to_decision_update)
erdos.add_watermark_callback([obstacles_stream, camera_stream],
[obstacle_tracking_stream],
self.on_watermark)
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._tracker_type = tracker_type
# Absolute time when the last tracker run completed.
self._last_tracker_run_completion_time = 0
try:
if tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(
self._flags, self._logger)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags,
self._logger)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError as error:
self._logger.fatal('Error importing {}'.format(tracker_type))
raise error
self._obstacles_msgs = deque()
self._frame_msgs = deque()
self._detection_update_count = -1
def __init__(self, obstacles_stream, camera_stream,
obstacle_tracking_stream, tracker_type, flags):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
erdos.add_watermark_callback([obstacles_stream, camera_stream],
[obstacle_tracking_stream],
self.on_watermark)
self._flags = flags
self._logger = erdos.utils.setup_logging(self.config.name,
self.config.log_file_name)
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import\
MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
self._obstacles_msgs = deque()
self._frame_msgs = deque()
def __init__(self,
name,
output_stream_name,
tracker_type,
flags,
log_file_name=None,
csv_file_name=None):
super(ObjectTrackerOp, self).__init__(name)
self._flags = flags
self._logger = setup_logging(self.name, log_file_name)
self._csv_logger = setup_csv_logging(self.name + '-csv', csv_file_name)
self._output_stream_name = output_stream_name
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
self._logger.fatal(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
# True when the tracker is ready to update bboxes.
self._ready_to_update = False
self._ready_to_update_timestamp = None
self._to_process = deque()
self._lock = threading.Lock()
class ObjectTrackerOperator(erdos.Operator):
def __init__(self, obstacles_stream, camera_stream,
time_to_decision_stream, obstacle_tracking_stream,
tracker_type, flags):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
time_to_decision_stream.add_callback(self.on_time_to_decision_update)
erdos.add_watermark_callback([obstacles_stream, camera_stream],
[obstacle_tracking_stream],
self.on_watermark)
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._tracker_type = tracker_type
try:
if tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(
self._flags, self._logger)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags,
self._logger)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
self._obstacles_msgs = deque()
self._frame_msgs = deque()
self._watermark_msg_count = 0
@staticmethod
def connect(obstacles_stream, camera_stream, time_to_decision_stream):
obstacle_tracking_stream = erdos.WriteStream()
return [obstacle_tracking_stream]
def on_frame_msg(self, msg):
"""Invoked when a FrameMessage is received on the camera stream."""
self._logger.debug('@{}: {} received frame'.format(
msg.timestamp, self.config.name))
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
self._frame_msgs.append(msg)
def on_obstacles_msg(self, msg):
"""Invoked when obstacles are received on the stream."""
self._logger.debug('@{}: {} received {} obstacles'.format(
msg.timestamp, self.config.name, len(msg.obstacles)))
self._obstacles_msgs.append(msg)
def on_time_to_decision_update(self, msg):
self._logger.debug('@{}: {} received ttd update {}'.format(
msg.timestamp, self.config.name, msg))
@erdos.profile_method()
def on_watermark(self, timestamp, obstacle_tracking_stream):
self._logger.debug('@{}: received watermark'.format(timestamp))
frame_msg = self._frame_msgs.popleft()
camera_frame = frame_msg.frame
tracked_obstacles = []
self._watermark_msg_count += 1
if len(self._obstacles_msgs) > 0:
obstacles_msg = self._obstacles_msgs.popleft()
assert frame_msg.timestamp == obstacles_msg.timestamp
detected_obstacles = []
for obstacle in obstacles_msg.obstacles:
if obstacle.is_vehicle() or obstacle.is_person():
detected_obstacles.append(obstacle)
if (self._watermark_msg_count %
self._flags.track_every_nth_detection == 0):
# Reinitialize the tracker with new detections.
self._logger.debug(
'Restarting trackers at frame {}'.format(timestamp))
self._tracker.reinitialize(camera_frame, detected_obstacles)
self._logger.debug('Processing frame {}'.format(timestamp))
ok, tracked_obstacles = self._tracker.track(camera_frame)
if not ok:
self._logger.error(
'Tracker failed at timestamp {}'.format(timestamp))
obstacle_tracking_stream.send(
ObstaclesMessage(timestamp, tracked_obstacles, 0))
obstacle_tracking_stream.send(erdos.WatermarkMessage(timestamp))
class ObjectTrackerOperator(erdos.Operator):
def __init__(self, obstacles_stream, camera_stream,
time_to_decision_stream, obstacle_tracking_stream,
tracker_type, flags):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
time_to_decision_stream.add_callback(self.on_time_to_decision_update)
erdos.add_watermark_callback([obstacles_stream, camera_stream],
[obstacle_tracking_stream],
self.on_watermark)
self._flags = flags
self._logger = erdos.utils.setup_logging(self.config.name,
self.config.log_file_name)
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import\
MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
self._obstacles_msgs = deque()
self._frame_msgs = deque()
@staticmethod
def connect(obstacles_stream, camera_stream, time_to_decision_stream):
obstacle_tracking_stream = erdos.WriteStream()
return [obstacle_tracking_stream]
def on_frame_msg(self, msg):
"""Invoked when a FrameMessage is received on the camera stream."""
self._logger.debug('@{}: {} received frame'.format(
msg.timestamp, self.config.name))
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
self._frame_msgs.append(msg)
def on_obstacles_msg(self, msg):
"""Invoked when obstacles are received on the stream."""
self._logger.debug('@{}: {} received {} obstacles'.format(
msg.timestamp, self.config.name, len(msg.obstacles)))
self._obstacles_msgs.append(msg)
def on_time_to_decision_update(self, msg):
self._logger.debug('@{}: {} received ttd update {}'.format(
msg.timestamp, self.config.name, msg))
@erdos.profile_method()
def on_watermark(self, timestamp, obstacle_tracking_stream):
self._logger.debug('@{}: received watermark'.format(timestamp))
frame_msg = self._frame_msgs.popleft()
camera_frame = frame_msg.frame
tracked_obstacles = []
if len(self._obstacles_msgs) > 0:
obstacles_msg = self._obstacles_msgs.popleft()
assert frame_msg.timestamp == obstacles_msg.timestamp
self._logger.debug(
'Restarting trackers at frame {}'.format(timestamp))
detected_obstacles = []
for obstacle in obstacles_msg.obstacles:
if (obstacle.label in VEHICLE_LABELS
or obstacle.label == 'person'):
detected_obstacles.append(obstacle)
self._tracker.reinitialize(camera_frame, detected_obstacles)
self._logger.debug('Processing frame {}'.format(timestamp))
ok, tracked_obstacles = self._tracker.track(camera_frame)
if not ok:
self._logger.error(
'Tracker failed at timestamp {}'.format(timestamp))
obstacle_tracking_stream.send(
ObstaclesMessage(timestamp, tracked_obstacles, 0))
if self._flags.visualize_tracker_output:
# Tracked obstacles have no label, draw white bbox.
camera_frame.annotate_with_bounding_boxes(timestamp,
tracked_obstacles)
camera_frame.visualize(self.config.name,
pygame_display=pylot.utils.PYGAME_DISPLAY)
class ObjectTrackerOperator(erdos.Operator):
def __init__(self, obstacles_stream, camera_stream,
time_to_decision_stream, obstacle_tracking_stream,
tracker_type, flags):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
time_to_decision_stream.add_callback(self.on_time_to_decision_update)
erdos.add_watermark_callback([obstacles_stream, camera_stream],
[obstacle_tracking_stream],
self.on_watermark)
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._tracker_type = tracker_type
# Absolute time when the last tracker run completed.
self._last_tracker_run_completion_time = 0
try:
if tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(
self._flags, self._logger)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags,
self._logger)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError as error:
self._logger.fatal('Error importing {}'.format(tracker_type))
raise error
self._obstacles_msgs = deque()
self._frame_msgs = deque()
self._detection_update_count = -1
@staticmethod
def connect(obstacles_stream, camera_stream, time_to_decision_stream):
obstacle_tracking_stream = erdos.WriteStream()
return [obstacle_tracking_stream]
def destroy(self):
self._logger.warn('destroying {}'.format(self.config.name))
def on_frame_msg(self, msg):
"""Invoked when a FrameMessage is received on the camera stream."""
self._logger.debug('@{}: {} received frame'.format(
msg.timestamp, self.config.name))
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
self._frame_msgs.append(msg)
def on_obstacles_msg(self, msg):
"""Invoked when obstacles are received on the stream."""
self._logger.debug('@{}: {} received {} obstacles'.format(
msg.timestamp, self.config.name, len(msg.obstacles)))
self._obstacles_msgs.append(msg)
def on_time_to_decision_update(self, msg):
self._logger.debug('@{}: {} received ttd update {}'.format(
msg.timestamp, self.config.name, msg))
def _reinit_tracker(self, camera_frame, detected_obstacles):
start = time.time()
result = self._tracker.reinitialize(camera_frame, detected_obstacles)
return (time.time() - start) * 1000, result
def _run_tracker(self, camera_frame):
start = time.time()
result = self._tracker.track(camera_frame)
return (time.time() - start) * 1000, result
@erdos.profile_method()
def on_watermark(self, timestamp, obstacle_tracking_stream):
self._logger.debug('@{}: received watermark'.format(timestamp))
if timestamp.is_top:
return
frame_msg = self._frame_msgs.popleft()
camera_frame = frame_msg.frame
tracked_obstacles = []
detector_runtime = 0
reinit_runtime = 0
# Check if the most recent obstacle message has this timestamp.
# If it doesn't, then the detector might have skipped sending
# an obstacle message.
if (len(self._obstacles_msgs) > 0
and self._obstacles_msgs[0].timestamp == timestamp):
obstacles_msg = self._obstacles_msgs.popleft()
self._detection_update_count += 1
if (self._detection_update_count %
self._flags.track_every_nth_detection == 0):
# Reinitialize the tracker with new detections.
self._logger.debug(
'Restarting trackers at frame {}'.format(timestamp))
detected_obstacles = []
for obstacle in obstacles_msg.obstacles:
if obstacle.is_vehicle() or obstacle.is_person():
detected_obstacles.append(obstacle)
reinit_runtime, _ = self._reinit_tracker(
camera_frame, detected_obstacles)
detector_runtime = obstacles_msg.runtime
tracker_runtime, (ok, tracked_obstacles) = \
self._run_tracker(camera_frame)
assert ok, 'Tracker failed at timestamp {}'.format(timestamp)
tracker_runtime = tracker_runtime + reinit_runtime
tracker_delay = self.__compute_tracker_delay(timestamp.coordinates[0],
detector_runtime,
tracker_runtime)
obstacle_tracking_stream.send(
ObstaclesMessage(timestamp, tracked_obstacles, tracker_delay))
def __compute_tracker_delay(self, world_time, detector_runtime,
tracker_runtime):
# If the tracker runtime does not fit within the frame gap, then
# the tracker will fall behind. We need a scheduler to better
# handle such situations.
if (world_time + detector_runtime >
self._last_tracker_run_completion_time):
# The detector finished after the previous tracker invocation
# completed. Therefore, the tracker is already sequenced.
tracker_runtime = detector_runtime + tracker_runtime
self._last_tracker_run_completion_time = \
world_time + tracker_runtime
else:
# The detector finished before the previous tracker invocation
# completed. The tracker can only run after the previous
# invocation completes.
self._last_tracker_run_completion_time += tracker_runtime
tracker_runtime = \
self._last_tracker_run_completion_time - world_time
return tracker_runtime
class ObjectTrackerOperator(erdos.Operator):
def __init__(self,
obstacles_stream,
camera_stream,
obstacle_tracking_stream,
name,
tracker_type,
flags,
log_file_name=None,
csv_file_name=None):
obstacles_stream.add_callback(self.on_obstacles_msg)
camera_stream.add_callback(self.on_frame_msg)
self._name = name
self._flags = flags
self._logger = erdos.utils.setup_logging(name, log_file_name)
self._csv_logger = erdos.utils.setup_csv_logging(
name + '-csv', csv_file_name)
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import\
MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import\
MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import\
MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import\
MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
raise ValueError(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
# Labels the obstacle trackers should track.
self._tracked_labels = {
'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck'
}
# True when the tracker is ready to update bboxes.
self._ready_to_update = False
self._ready_to_update_timestamp = None
self._to_process = deque()
@staticmethod
def connect(obstacles_stream, camera_stream):
obstacle_tracking_stream = erdos.WriteStream()
return [obstacle_tracking_stream]
def on_frame_msg(self, msg):
""" Invoked when a FrameMessage is received on the camera stream."""
self._logger.debug('@{}: {} received frame'.format(
msg.timestamp, self._name))
start_time = time.time()
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
camera_frame = msg.frame
# Store frames so that they can be re-processed once we receive the
# next update from the detector.
self._to_process.append((msg.timestamp, camera_frame))
# Track if we have a tracker ready to accept new frames.
if self._ready_to_update:
self.__track_bboxes_on_frame(camera_frame, msg.timestamp, False)
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
self._csv_logger.info('{},{},"{}",{}'.format(time_epoch_ms(),
self._name, msg.timestamp,
runtime))
def on_obstacles_msg(self, msg):
""" Invoked when obstacles are received on the stream."""
self._logger.debug('@{}: {} received obstacles'.format(
msg.timestamp, self._name))
self._ready_to_update = False
self._logger.debug("@{}: received {} bounding boxes".format(
msg.timestamp, len(msg.obstacles)))
# Remove frames that are older than the detector update.
while len(self._to_process
) > 0 and self._to_process[0][0] < msg.timestamp:
self._logger.debug("@{}: removing stale {} {}".format(
msg.timestamp, self._to_process[0][0], msg.timestamp))
self._to_process.popleft()
tracked_obstacles = []
for obstacle in msg.obstacles:
if obstacle.label in self._tracked_labels:
tracked_obstacles.append(obstacle)
if len(tracked_obstacles) > 0:
if len(self._to_process) > 0:
# Found the frame corresponding to the bounding boxes.
(timestamp, camera_frame) = self._to_process.popleft()
assert timestamp == msg.timestamp
# Re-initialize trackers.
self._ready_to_update = True
self._ready_to_update_timestamp = timestamp
self._logger.debug(
'Restarting trackers at frame {}'.format(timestamp))
self._tracker.reinitialize(camera_frame, tracked_obstacles)
self._logger.debug(
'Trackers have {} frames to catch-up'.format(
len(self._to_process)))
for (timestamp, camera_frame) in self._to_process:
if self._ready_to_update:
self.__track_bboxes_on_frame(camera_frame, timestamp,
True)
else:
self._logger.debug(
'@{}: received bboxes update, but no frame to process'.
format(msg.timestamp))
def __track_bboxes_on_frame(self, camera_frame, timestamp, catch_up):
self._logger.debug('Processing frame {}'.format(timestamp))
# Sequentually update state for each bounding box.
ok, tracked_obstacles = self._tracker.track(camera_frame)
if not ok:
self._logger.error(
'Tracker failed at timestamp {} last ready_to_update at {}'.
format(timestamp, self._ready_to_update_timestamp))
# The tracker must be reinitialized.
self._ready_to_update = False
else:
if self._flags.visualize_tracker_output and not catch_up:
# tracked obstacles have no label, draw white bbox.
camera_frame.annotate_with_bounding_boxes(
timestamp, tracked_obstacles)
camera_frame.visualize(self._name)
class ObjectTrackerOp(Op):
def __init__(self,
name,
output_stream_name,
tracker_type,
flags,
log_file_name=None,
csv_file_name=None):
super(ObjectTrackerOp, self).__init__(name)
self._flags = flags
self._logger = setup_logging(self.name, log_file_name)
self._csv_logger = setup_csv_logging(self.name + '-csv', csv_file_name)
self._output_stream_name = output_stream_name
self._tracker_type = tracker_type
try:
if tracker_type == 'cv2':
from pylot.perception.tracking.cv2_tracker import MultiObjectCV2Tracker
self._tracker = MultiObjectCV2Tracker(self._flags)
elif tracker_type == 'da_siam_rpn':
from pylot.perception.tracking.da_siam_rpn_tracker import MultiObjectDaSiamRPNTracker
self._tracker = MultiObjectDaSiamRPNTracker(self._flags)
elif tracker_type == 'deep_sort':
from pylot.perception.tracking.deep_sort_tracker import MultiObjectDeepSORTTracker
self._tracker = MultiObjectDeepSORTTracker(
self._flags, self._logger)
elif tracker_type == 'sort':
from pylot.perception.tracking.sort_tracker import MultiObjectSORTTracker
self._tracker = MultiObjectSORTTracker(self._flags)
else:
self._logger.fatal(
'Unexpected tracker type {}'.format(tracker_type))
except ImportError:
self._logger.fatal('Error importing {}'.format(tracker_type))
# True when the tracker is ready to update bboxes.
self._ready_to_update = False
self._ready_to_update_timestamp = None
self._to_process = deque()
self._lock = threading.Lock()
@staticmethod
def setup_streams(input_streams,
output_stream_name,
camera_stream_name=None):
input_streams.filter(is_obstacles_stream).add_callback(
ObjectTrackerOp.on_objects_msg)
# Select camera input streams.
camera_streams = input_streams.filter(is_camera_stream)
if camera_stream_name:
# Select only the camera the operator is interested in.
camera_streams = camera_streams.filter_name(camera_stream_name)
# Register a callback on the camera input stream.
camera_streams.add_callback(ObjectTrackerOp.on_frame_msg)
return [DataStream(name=output_stream_name)]
def on_frame_msg(self, msg):
""" Invoked when a FrameMessage is received on the camera stream."""
self._lock.acquire()
start_time = time.time()
assert msg.encoding == 'BGR', 'Expects BGR frames'
frame = msg.frame
# Store frames so that they can be re-processed once we receive the
# next update from the detector.
self._to_process.append((msg.timestamp, frame))
# Track if we have a tracker ready to accept new frames.
if self._ready_to_update:
self.__track_bboxes_on_frame(frame, msg.timestamp, False)
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
self._csv_logger.info('{},{},"{}",{}'.format(time_epoch_ms(),
self.name, msg.timestamp,
runtime))
self._lock.release()
def on_objects_msg(self, msg):
""" Invoked when detected objects are received on the stream."""
self._lock.acquire()
self._ready_to_update = False
self._logger.info("Received {} bboxes for {}".format(
len(msg.detected_objects), msg.timestamp))
# Remove frames that are older than the detector update.
while len(self._to_process
) > 0 and self._to_process[0][0] < msg.timestamp:
self._logger.info("Removing stale {} {}".format(
self._to_process[0][0], msg.timestamp))
self._to_process.popleft()
# Track all pedestrians.
bboxes, ids, confidence_scores = self.__get_pedestrians(
msg.detected_objects) # xmin, ymin, xmax, ymax
if len(bboxes) > 0:
if len(self._to_process) > 0:
# Found the frame corresponding to the bounding boxes.
(timestamp, frame) = self._to_process.popleft()
assert timestamp == msg.timestamp
# Re-initialize trackers.
self.__initialize_trackers(frame, bboxes, msg.timestamp,
confidence_scores)
self._logger.info('Trackers have {} frames to catch-up'.format(
len(self._to_process)))
for (timestamp, frame) in self._to_process:
if self._ready_to_update:
self.__track_bboxes_on_frame(frame, timestamp, True)
else:
self._logger.info(
'Received bboxes update {}, but no frame to process'.
format(msg.timestamp))
self._lock.release()
def execute(self):
self.spin()
def checkpoint(self, timestamp):
""" Invoked by ERDOS when the operator must checkpoint state.
The checkpoint should include all the state up to and including
timestamp.
Args:
timestamp: The timestamp at which the checkpoint is taken.
"""
# We can't checkpoint the tracker itself because it has internal state.
state = [
self._ready_to_update_timestamp, self._ready_to_update,
self._to_process
]
file_name = '{}{}.checkpoint'.format(self._name,
timestamp.coordinates[0])
pickle.dump(state, open(file_name, 'wb'))
return file_name
def restore(self, timestamp, checkpoint):
""" Invoked by ERDOS when the operator must restore state.
The method rebuilds the state from the checkpoint.
Args:
timestamp: The timestamp the checkpoint was taken at.
checkpoint: The checkpoint was taken at timestamp.
"""
checkpoint = pickle.load(open(checkpoint, 'rb'))
self._ready_to_update_timestamp = checkpoint[0]
self._ready_to_update = checkpoint[1]
self._to_process = checkpoint[2]
def __get_highest_confidence_pedestrian(self, detected_objs):
max_confidence = 0
max_corners = None
for detected_obj in detected_objs:
if (detected_obj.label == 'person'
and detected_obj.confidence > max_confidence):
max_corners = detected_obj.corners
max_confidence = detected_obj.confidence
if max_corners:
return [max_corners]
else:
return []
def __get_pedestrians(self, detector_objs):
bboxes = []
ids = []
confidence_scores = []
for detected_obj in detector_objs:
if detected_obj.label == 'person':
bboxes.append(detected_obj.corners)
ids.append(detected_obj.obj_id)
confidence_scores.append(detected_obj.confidence)
return bboxes, ids, confidence_scores
def __initialize_trackers(self, frame, bboxes, timestamp,
confidence_scores):
self._ready_to_update = True
self._ready_to_update_timestamp = timestamp
self._logger.info('Restarting trackers at frame {}'.format(timestamp))
self._tracker.reinitialize(frame, bboxes, confidence_scores)
def __track_bboxes_on_frame(self, frame, timestamp, catch_up):
self._logger.info('Processing frame {}'.format(timestamp))
# Sequentually update state for each bounding box.
ok, tracked_objects = self._tracker.track(frame)
if not ok:
self._logger.error(
'Tracker failed at timestamp {} last ready_to_update at {}'.
format(timestamp, self._ready_to_update_timestamp))
# The tracker must be reinitialized.
self._ready_to_update = False
else:
if self._flags.visualize_tracker_output and not catch_up:
for tracked_object in tracked_objects:
# tracked objects have no label, draw black bbox for them ([0, 0, 0])
tracked_object.visualize_on_img(frame, {"": [0, 0, 0]})
visualize_image(self.name, frame)