def read_geometry(self, run, tag, step, batch_idx, step_to_idx):
"""Geometry reader from msgpack files."""
idx = step_to_idx[step]
metadata_proto = plugin_data_pb2.Open3DPluginData()
run_tensor_events = self.tensor_events(run)
metadata_proto.ParseFromString(
run_tensor_events[tag][idx].tensor_proto.string_val[0])
data_dir = PluginDirectory(os.path.join(self.logdir, run),
metadata.PLUGIN_NAME)
filename = os.path.join(data_dir, metadata_proto.batch_index.filename)
read_location = metadata_proto.batch_index.start_size[batch_idx].start
read_size = metadata_proto.batch_index.start_size[batch_idx].size
read_masked_crc32c = metadata_proto.batch_index.start_size[
batch_idx].masked_crc32c
cache_key = (filename, read_location, read_size, run, tag, step,
batch_idx)
geometry = self.geometry_cache.get(cache_key)
if geometry is None: # Read from storage
buf = self.read_from_file(filename, read_location, read_size,
read_masked_crc32c)
if buf is None:
raise IOError(f"Geometry {cache_key} reading failed! CRC "
"mismatch in msgpack data.")
msg_tag, msg_step, geometry = o3d.io.rpc.data_buffer_to_meta_geometry(
buf)
if geometry is None:
raise IOError(f"Geometry {cache_key} reading failed! Possible "
"msgpack or TensorFlow event file corruption.")
if tag != msg_tag or step != msg_step:
_log.warning(
f"Mismatch between TensorFlow event (tag={tag}, step={step})"
f" and msgpack (tag={msg_tag}, step={msg_step}) data. "
"Possible data corruption.")
_log.debug(f"Geometry {cache_key} reading successful!")
self.geometry_cache.put(cache_key, geometry)
# Fill in properties by reference
for prop_ref in metadata_proto.property_references:
prop = plugin_data_pb2.Open3DPluginData.GeometryProperty.Name(
prop_ref.geometry_property)
if prop_ref.step_ref >= step:
_log.warning(
f"Incorrect future step reference {prop_ref.step_ref} for"
f" property {prop} of geometry at step {step}. Ignoring.")
continue
geometry_ref = self.read_geometry(run, tag, prop_ref.step_ref,
batch_idx, step_to_idx)[0]
# "vertex_normals" -> ["vertex", "normals"]
prop_map, prop_attribute = prop.split("_")
if prop_map == "vertex" and not isinstance(
geometry, o3d.t.geometry.TriangleMesh):
prop_map = "point"
# geometry.vertex["normals"] = geometry_ref.vertex["normals"]
getattr(geometry, prop_map)[prop_attribute] = getattr(
geometry_ref, prop_map)[prop_attribute]
# Aux data (e.g. labels, confidences) -> not cached
aux_read_location = metadata_proto.batch_index.start_size[
batch_idx].aux_start
aux_read_size = metadata_proto.batch_index.start_size[
batch_idx].aux_size
aux_read_masked_crc32c = metadata_proto.batch_index.start_size[
batch_idx].aux_masked_crc32c
data_bbox_proto = plugin_data_pb2.InferenceData()
if aux_read_size > 0:
data_bbox_serial = self.read_from_file(filename, aux_read_location,
aux_read_size,
aux_read_masked_crc32c)
if data_bbox_serial is None:
raise IOError(f"Aux data for {cache_key} reading failed! CRC "
"mismatch in protobuf data.")
data_bbox_proto.ParseFromString(data_bbox_serial)
self.update_runtag_prop_shape(run, tag, geometry, data_bbox_proto)
return geometry, data_bbox_proto
def read_geometry(self, run, tag, step, batch_idx, step_to_idx):
"""Geometry reader from msgpack files.
"""
idx = step_to_idx[step]
metadata_proto = plugin_data_pb2.Open3DPluginData()
run_tensor_events = self.tensor_events(run)
metadata_proto.ParseFromString(
run_tensor_events[tag][idx].tensor_proto.string_val[0])
data_dir = PluginDirectory(os.path.join(self.logdir, run),
metadata.PLUGIN_NAME)
filename = os.path.join(data_dir, metadata_proto.batch_index.filename)
read_location = metadata_proto.batch_index.start_size[batch_idx].start
read_size = metadata_proto.batch_index.start_size[batch_idx].size
read_masked_crc32c = metadata_proto.batch_index.start_size[
batch_idx].masked_crc32c
cache_key = (filename, read_location, read_size, run, tag, step,
batch_idx)
geometry = self.geometry_cache.get(cache_key)
if geometry is None: # Read from storage
with self._file_handles_lock:
if filename not in self._file_handles:
self._file_handles[filename] = (_fileopen(filename, "rb"),
threading.Lock())
if not self._file_handles[filename][0].seekable():
raise RuntimeError(filename +
" does not support seeking."
" This storage is not supported.")
# lock to seek + read
file_handle = self._file_handles[filename]
file_handle[1].acquire()
file_handle[0].seek(read_location)
buf = file_handle[0].read(read_size)
file_handle[1].release()
if not read_masked_crc32c == masked_crc32c(buf):
raise IOError(f"Geometry {cache_key} reading failed! CRC "
"mismatch in msgpack data.")
msg_tag, msg_step, geometry = o3d.io.rpc.data_buffer_to_meta_geometry(
buf)
if geometry is None:
raise IOError(f"Geometry {cache_key} reading failed! Possible "
"msgpack or TensorFlow event file corruption.")
if tag != msg_tag or step != msg_step:
_log.warning(
f"Mismatch between TensorFlow event (tag={tag}, step={step})"
f" and msgpack (tag={msg_tag}, step={msg_step}) data. "
"Possible data corruption.")
_log.debug(f"Geometry {cache_key} reading successful!")
self.geometry_cache.put(cache_key, geometry)
# Fill in properties by reference
for prop_ref in metadata_proto.property_references:
prop = plugin_data_pb2.Open3DPluginData.GeometryProperty.Name(
prop_ref.geometry_property)
if prop_ref.step_ref >= step:
_log.warning(
f"Incorrect future step reference {prop_ref.step_ref} for"
f" property {prop} of geometry at step {step}. Ignoring.")
continue
geometry_ref = self.read_geometry(run, tag, prop_ref.step_ref,
batch_idx, step_to_idx)
# "vertex_normals" -> ["vertex", "normals"]
prop_map, prop_attribute = prop.split("_")
if prop_map == "vertex" and not isinstance(
geometry, o3d.t.geometry.TriangleMesh):
prop_map = "point"
# geometry.vertex["normals"] = geometry_ref.vertex["normals"]
getattr(geometry, prop_map)[prop_attribute] = getattr(
geometry_ref, prop_map)[prop_attribute]
return geometry
def lazy_tensor():
write_dir = PluginDirectory(logdir, metadata.PLUGIN_NAME)
geometry_metadata_string = _write_geometry_data(
write_dir, tag, step, data, max_outputs)
return tf.convert_to_tensor(geometry_metadata_string)