def enqueue(self, tagfilepath, data):
"""Add a write job to the write queue.
Args:
tagfilepath (str): Full file pathname for data. A suffix will be
added to get the complete filename. A None value indicates
writing is over and the writer thread should join.
data (bytes): Data buffer to write.
Returns:
Tuple of filename and location (in bytes) where the data will be
written.
"""
with self._file_lock:
if tagfilepath not in self._file_handles:
# summary.writer.EventFileWriter name format
fullfilepath = "{}.{}.{}.{}{}".format(tagfilepath,
int(time.time()),
socket.gethostname(),
os.getpid(),
self._filename_extension)
_makedirs(os.path.dirname(fullfilepath))
self._file_handles[tagfilepath] = _fileopen(fullfilepath, 'wb')
_log.debug(f"msgpack file {fullfilepath} opened for writing.")
this_write_loc = 0
self._file_next_write_pos[tagfilepath] = len(data)
else:
this_write_loc = self._file_next_write_pos[tagfilepath]
self._file_next_write_pos[tagfilepath] += len(data)
fullfilepath = self._file_handles[tagfilepath].name
# Blocks till queue has available slot.
self._write_queue.put((tagfilepath, data), block=True)
if not self._writer_thread.is_alive():
self._writer_thread = threading.Thread(target=self._writer,
name="Open3DDataWriter")
self._writer_thread.start()
return os.path.basename(fullfilepath), this_write_loc
def read_from_file(self, filename, read_location, read_size,
read_masked_crc32c):
"""Read data from the file ``filename`` from a given offset
``read_location`` and size ``read_size``. Data is validated with the provided
``masked_crc32c``. This is thread safe and manages a list of open files.
"""
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 masked_crc32c(buf) == read_masked_crc32c:
return buf
else:
return None
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
