def __init__(self,
camera_stream,
segmented_stream,
name,
flags,
log_file_name=None,
csv_file_name=None):
camera_stream.add_callback(self.on_msg_camera_stream,
[segmented_stream])
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)
arch = "drn_d_22"
classes = 19
self._pallete = drn.segment.CARLA_CITYSCAPE_PALETTE
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(
torch.load(self._flags.segmentation_model_path))
if torch.cuda.is_available():
self._model = torch.nn.DataParallel(self._model).cuda()
def __init__(self,
name,
output_stream_name,
flags,
log_file_name=None,
csv_file_name=None):
super(SegmentationDRNOperator, 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
arch = "drn_d_22"
classes = 19
pretrained = "dependencies/data/drn_d_22_cityscapes.pth"
self._pallete = drn.segment.CITYSCAPE_PALETTE
# TODO(ionel): Figure out how to set GPU memory fraction.
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(torch.load(pretrained))
# TODO(ionel): Automatically detect if GPU is available.
if self._flags.segmentation_gpu:
self._model = torch.nn.DataParallel(self._model).cuda()
self._last_seq_num = -1
def __init__(self,
name,
output_stream_name,
flags,
log_file_name=None,
csv_file_name=None):
super(SegmentationDRNOperator, 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
arch = "drn_d_22"
classes = 19
self._pallete = drn.segment.CITYSCAPE_PALETTE
# TODO(ionel): Figure out how to set GPU memory fraction.
self._model = DRNSeg(
arch, classes, pretrained_model=None, pretrained=False)
self._model.load_state_dict(
torch.load(self._flags.segmentation_drn_model_path))
if torch.cuda.is_available():
self._model = torch.nn.DataParallel(self._model).cuda()
class SegmentationDRNOperator(erdos.Operator):
"""Semantically segments frames using a DRN segmentation model.
The operator receives frames on a camera stream, and runs a model for each
frame.
Args:
camera_stream (:py:class:`erdos.ReadStream`): The stream on which
camera frames are received.
segmented_stream (:py:class:`erdos.WriteStream`): Stream on which the
operator sends
:py:class:`~pylot.perception.messages.SegmentedFrameMessage`
messages.
flags (absl.flags): Object to be used to access absl flags.
"""
def __init__(self, camera_stream, segmented_stream, flags):
camera_stream.add_callback(self.on_msg_camera_stream,
[segmented_stream])
self._flags = flags
self._logger = erdos.utils.setup_logging(self.config.name,
self.config.log_file_name)
arch = "drn_d_22"
classes = 19
self._pallete = drn.segment.CARLA_CITYSCAPE_PALETTE
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(
torch.load(self._flags.segmentation_model_path))
if torch.cuda.is_available():
self._model = torch.nn.DataParallel(self._model).cuda()
@staticmethod
def connect(camera_stream):
"""Connects the operator to other streams.
Args:
camera_stream (:py:class:`erdos.ReadStream`): The stream on which
camera frames are received.
Returns:
:py:class:`erdos.WriteStream`: Stream on which the operator sends
:py:class:`~pylot.perception.messages.SegmentedFrameMessage`
messages.
"""
segmented_stream = erdos.WriteStream()
return [segmented_stream]
@erdos.profile_method()
def on_msg_camera_stream(self, msg, segmented_stream):
"""Invoked upon the receipt of a message on the camera stream.
Args:
msg: A :py:class:`~pylot.perception.messages.FrameMessage`.
segmented_stream (:py:class:`erdos.WriteStream`): Stream on which
the operator sends
:py:class:`~pylot.perception.messages.SegmentedFrameMessage`
messages.
"""
self._logger.debug('@{}: {} received message'.format(
msg.timestamp, self.config.name))
start_time = time.time()
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
image = torch.from_numpy(msg.frame.frame.transpose(
[2, 0, 1])).unsqueeze(0).float()
image_var = Variable(image, requires_grad=False, volatile=True)
final = self._model(image_var)[0]
_, pred = torch.max(final, 1)
pred = pred.cpu().data.numpy()[0]
# After we apply the pallete, the image is in RGB format
image_np = self._pallete[pred.squeeze()]
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
frame = SegmentedFrame(image_np, 'cityscapes', msg.frame.camera_setup)
if self._flags.visualize_segmentation_output:
frame.visualize(self.config.name,
msg.timestamp,
pygame_display=pylot.utils.PYGAME_DISPLAY)
segmented_stream.send(
SegmentedFrameMessage(msg.timestamp, frame, runtime))
class SegmentationDRNOperator(erdos.Operator):
""" Subscribes to a camera stream, and segments frames using DRN."""
def __init__(self,
camera_stream,
segmented_stream,
name,
flags,
log_file_name=None,
csv_file_name=None):
camera_stream.add_callback(self.on_msg_camera_stream,
[segmented_stream])
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)
arch = "drn_d_22"
classes = 19
self._pallete = drn.segment.CARLA_CITYSCAPE_PALETTE
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(
torch.load(self._flags.segmentation_model_path))
if torch.cuda.is_available():
self._model = torch.nn.DataParallel(self._model).cuda()
@staticmethod
def connect(camera_stream):
segmented_stream = erdos.WriteStream()
return [segmented_stream]
def on_msg_camera_stream(self, msg, segmented_stream):
"""Camera stream callback method.
Invoked upon the receipt of a message on the camera stream.
"""
self._logger.debug('@{}: {} received message'.format(
msg.timestamp, self._name))
start_time = time.time()
assert msg.frame.encoding == 'BGR', 'Expects BGR frames'
image = torch.from_numpy(msg.frame.frame.transpose(
[2, 0, 1])).unsqueeze(0).float()
image_var = Variable(image, requires_grad=False, volatile=True)
final = self._model(image_var)[0]
_, pred = torch.max(final, 1)
pred = pred.cpu().data.numpy()[0]
# After we apply the pallete, the image is in RGB format
image_np = self._pallete[pred.squeeze()]
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
self._csv_logger.info('{},{},"{}",{}'.format(time_epoch_ms(),
self._name, msg.timestamp,
runtime))
frame = SegmentedFrame(image_np, 'cityscapes', msg.frame.camera_setup)
if self._flags.visualize_segmentation_output:
frame.visualize(self._name, msg.timestamp)
segmented_stream.send(
SegmentedFrameMessage(frame, msg.timestamp, runtime))
class SegmentationDRNOperator(Op):
def __init__(self,
name,
output_stream_name,
flags,
log_file_name=None,
csv_file_name=None):
super(SegmentationDRNOperator, 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
arch = "drn_d_22"
classes = 19
pretrained = "dependencies/data/drn_d_22_cityscapes.pth"
self._pallete = drn.segment.CITYSCAPE_PALETTE
# TODO(ionel): Figure out how to set GPU memory fraction.
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(torch.load(pretrained))
# TODO(ionel): Automatically detect if GPU is available.
if self._flags.segmentation_gpu:
self._model = torch.nn.DataParallel(self._model).cuda()
self._last_seq_num = -1
@staticmethod
def setup_streams(input_streams, output_stream_name):
# Register a callback on the camera input stream.
input_streams.filter(is_camera_stream).add_callback(
SegmentationDRNOperator.on_msg_camera_stream)
return [create_segmented_camera_stream(output_stream_name)]
def on_msg_camera_stream(self, msg):
"""Camera stream callback method.
Invoked upon the receipt of a message on the camera stream.
"""
if self._last_seq_num + 1 != msg.timestamp.coordinates[1]:
self._logger.error(
'Expected msg with seq num {} but received {}'.format(
(self._last_seq_num + 1), msg.timestamp.coordinates[1]))
if self._flags.fail_on_message_loss:
assert self._last_seq_num + 1 == msg.timestamp.coordinates[1]
self._last_seq_num = msg.timestamp.coordinates[1]
self._logger.info('{} received frame {}'.format(
self.name, msg.timestamp))
start_time = time.time()
image = bgra_to_bgr(msg.data)
image = torch.from_numpy(image.transpose([2, 0,
1])).unsqueeze(0).float()
image_var = Variable(image, requires_grad=False, volatile=True)
final = self._model(image_var)[0]
_, pred = torch.max(final, 1)
pred = pred.cpu().data.numpy()[0]
image_np = self._pallete[pred.squeeze()]
# After we apply the pallete, the image is in RGB format
image_np = rgb_to_bgr(image_np)
if self._flags.visualize_segmentation_output:
add_timestamp(msg.timestamp, image_np)
cv2.imshow(self.name, image_np)
cv2.waitKey(1)
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
self._csv_logger.info('{},{},"{}",{}'.format(time_epoch_ms(),
self.name, msg.timestamp,
runtime))
output_msg = Message((image_np, runtime), msg.timestamp)
self.get_output_stream(self._output_stream_name).send(output_msg)
def execute(self):
"""Operator execute entry method."""
# Ensures that the operator runs continuously.
self.spin()
class SegmentationDRNOperator(Op):
""" Subscribes to a camera stream, and segments frames using DRN."""
def __init__(self,
name,
output_stream_name,
flags,
log_file_name=None,
csv_file_name=None):
super(SegmentationDRNOperator, 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
arch = "drn_d_22"
classes = 19
self._pallete = drn.segment.CITYSCAPE_PALETTE
# TODO(ionel): Figure out how to set GPU memory fraction.
self._model = DRNSeg(arch,
classes,
pretrained_model=None,
pretrained=False)
self._model.load_state_dict(
torch.load(self._flags.segmentation_drn_model_path))
# TODO(ionel): Automatically detect if GPU is available.
if self._flags.segmentation_gpu:
self._model = torch.nn.DataParallel(self._model).cuda()
@staticmethod
def setup_streams(input_streams,
output_stream_name,
filter_stream_name=None):
# Select camera input streams.
camera_streams = input_streams.filter(is_camera_stream)
if filter_stream_name:
# Select only the camera the operator is interested in.
camera_streams = camera_streams.filter_name(filter_stream_name)
# Register a callback on the camera input stream.
camera_streams.add_callback(
SegmentationDRNOperator.on_msg_camera_stream)
return [create_segmented_camera_stream(output_stream_name)]
def on_msg_camera_stream(self, msg):
"""Camera stream callback method.
Invoked upon the receipt of a message on the camera stream.
"""
self._logger.info('{} received frame {}'.format(
self.name, msg.timestamp))
start_time = time.time()
assert msg.encoding == 'BGR', 'Expects BGR frames'
image = torch.from_numpy(msg.frame.transpose([2, 0, 1
])).unsqueeze(0).float()
image_var = Variable(image, requires_grad=False, volatile=True)
final = self._model(image_var)[0]
_, pred = torch.max(final, 1)
pred = pred.cpu().data.numpy()[0]
image_np = self._pallete[pred.squeeze()]
# After we apply the pallete, the image is in RGB format
image_np = rgb_to_bgr(image_np)
if self._flags.visualize_segmentation_output:
add_timestamp(msg.timestamp, image_np)
cv2.imshow(self.name, image_np)
cv2.waitKey(1)
# Get runtime in ms.
runtime = (time.time() - start_time) * 1000
self._csv_logger.info('{},{},"{}",{}'.format(time_epoch_ms(),
self.name, msg.timestamp,
runtime))
output_msg = SegmentedFrameMessage(image_np, runtime, msg.timestamp)
self.get_output_stream(self._output_stream_name).send(output_msg)
def execute(self):
"""Operator execute entry method."""
# Ensures that the operator runs continuously.
self.spin()