def _validate_step(self, selected_step):
"""Validate step assuming self.run and self.tags are valid. Use
self.step or first valid step if selected_step is invalid."""
if len(self.tags) == 0: # No tags in this run
return
if selected_step not in self.step_to_idx:
selected_step = self.step
if selected_step not in self.step_to_idx:
selected_step = self.step_limits[0] # Set to first step
self.step = selected_step
self.idx = self.step_to_idx[self.step]
self.wall_time = self.all_tensor_events[self.tags[0]][
self.idx].wall_time
metadata_proto = plugin_data_pb2.Open3DPluginData()
metadata_proto.ParseFromString(self.all_tensor_events[self.tags[0]][
self.idx].tensor_proto.string_val[0])
self.batch_size = len(metadata_proto.batch_index.start_size)
def _write_geometry_data(write_dir, tag, step, data, max_outputs=1):
"""Serialize and write geometry data for a tag. Data is written to a
separate file per tag.
TODO: Add version specific reader / writer
Args:
write_dir (str): Path of folder to write data file.
tag (str): Full name for geometry.
step (int): Iteration / step count.
data (dict): Property name to tensor mapping.
max_outputs (int): Only the first `max_samples` data points in each
batch will be saved.
Returns:
A comma separated data location string with the format
f"{filename},{write_location},{write_size}"
"""
if not isinstance(data, dict):
raise TypeError(
"data should be a dict of geometry property names and tensors.")
data_bbox = None
if 'bboxes' in data:
if len(data) > 1:
raise ValueError("Saving bounding boxes. Please add other geometry"
" data with a separate call.")
data, data_bbox = _convert_bboxes(data['bboxes'])
unknown_props = [
prop for prop in data
if (prop not in metadata.SUPPORTED_PROPERTIES and
not prop.startswith('vertex_')) # custom vertex properties
]
if unknown_props:
raise ValueError(
f"Unknown geometry properties in data: {unknown_props}")
if "vertex_positions" not in data:
raise ValueError("Primary key 'vertex_positions' not provided.")
if max_outputs == 0: # save all
max_outputs = np.iinfo(np.int32).max
elif max_outputs < 1:
raise ValueError(
f"max_outputs ({max_outputs}) should be a non-negative integer.")
max_outputs = int(max_outputs)
if 'triangle_texture_uvs' in data and 'vertex_texture_uvs' in data:
raise ValueError("Please provide only one of triangle_texture_uvs "
"or vertex_texture_uvs.")
batch_size = None
n_vertices = None
n_triangles = None
n_lines = None
vertex_data = {}
triangle_data = {}
line_data = {}
material = {
"name": "",
"scalar_properties": {},
"vector_properties": {},
"texture_maps": {}
}
def get_or_check_shape(prop, tensor_tuple, exp_shape):
shape = [len(tensor_tuple)] + list(tensor_tuple[0].shape)
if len(shape) > 1:
shape[1] = tuple(tensor.shape[0] for tensor in tensor_tuple)
for k, s in enumerate(exp_shape):
if k < 2 and s is not None and s != shape[k]:
raise ValueError(f"Property {prop} tensor should be of shape "
f"{exp_shape} but is {shape}.")
if k > 1 and s is None:
s = tensor_tuple[0].shape[k - 1]
if k > 1 and any(
s != tensor.shape[k - 1] for tensor in tensor_tuple):
raise ValueError(
f"Property {prop} tensor should have shape[{k}]"
f"={s} for all elements but is "
f"{tensor.shape[k-1] for tensor in tensor_tuple}.")
return shape[:2]
geometry_metadata = plugin_data_pb2.Open3DPluginData(
version=metadata._VERSION)
o3d_type = "PointCloud"
for prop, tensor in data.items():
if prop.startswith('vertex_'):
prop_name = prop[7:]
vertex_data[prop_name] = _preprocess(prop, tensor, step,
max_outputs, geometry_metadata)
if vertex_data[prop_name] is None: # Step reference
del vertex_data[prop_name]
continue
# Get tensor dims from earlier property
batch_size, n_vertices = get_or_check_shape(
prop, vertex_data[prop_name], (batch_size, n_vertices) +
metadata.GEOMETRY_PROPERTY_DIMS.get(prop, (None,)))
elif prop in ('triangle_indices',) + metadata.TRIANGLE_PROPERTIES:
o3d_type = "TriangleMesh"
prop_name = prop[9:]
triangle_data[prop_name] = _preprocess(prop, tensor, step,
max_outputs,
geometry_metadata)
if triangle_data[prop_name] is None: # Step reference
del triangle_data[prop_name]
continue
# Get tensor dims from earlier property
batch_size, n_triangles = get_or_check_shape(
prop, triangle_data[prop_name], (batch_size, n_triangles) +
metadata.GEOMETRY_PROPERTY_DIMS.get(prop, (None,)))
elif prop in ('line_indices',) + metadata.LINE_PROPERTIES:
if o3d_type != "TriangleMesh":
o3d_type = "LineSet"
prop_name = prop[5:]
line_data[prop_name] = _preprocess(prop, tensor, step, max_outputs,
geometry_metadata)
if line_data[prop_name] is None: # Step reference
del line_data[prop_name]
continue
# Get tensor dims from earlier property
batch_size, n_lines = get_or_check_shape(
prop, line_data[prop_name], (batch_size, n_lines) +
metadata.GEOMETRY_PROPERTY_DIMS.get(prop, (None,)))
elif prop.startswith("material_"): # Set property in material directly
_preprocess(prop, tensor, step, max_outputs, geometry_metadata,
material)
if (len(material["texture_maps"]) > 0 and
'triangle_texture_uvs' not in data and
'vertex_texture_uvs' not in data):
raise ValueError("Please provide texture coordinates "
"'[vertex|triangle]_texture_uvs' to use texture maps.")
if material["name"] == "" and any(
len(value) > 0 for value in material.values()):
raise ValueError(
"Please provide a 'material_name' for the material properties.")
vertices = vertex_data.pop("positions",
o3d.core.Tensor((), dtype=o3d.core.float32))
faces = triangle_data.pop("indices",
o3d.core.Tensor((), dtype=o3d.core.int32))
lines = line_data.pop("indices", o3d.core.Tensor((), dtype=o3d.core.int32))
for bidx in range(batch_size):
buf_con = o3d.io.rpc.BufferConnection()
if not o3d.io.rpc.set_mesh_data(
path=tag,
time=step,
layer="",
vertices=vertices[bidx] if len(vertices) > 0 else vertices,
vertex_attributes={
prop: tensor[bidx] for prop, tensor in vertex_data.items()
},
faces=faces[bidx] if len(faces) > 0 else faces,
face_attributes={
prop: tensor[bidx]
for prop, tensor in triangle_data.items()
},
lines=lines[bidx] if len(lines) > 0 else lines,
line_attributes={
prop: tensor[bidx] for prop, tensor in line_data.items()
},
material=("" if material["name"] == "" else
material["name"][bidx]),
material_scalar_attributes={
prop: val[bidx].item()
for prop, val in material["scalar_properties"].items()
},
material_vector_attributes={
prop: val[bidx].numpy()
for prop, val in material["vector_properties"].items()
},
texture_maps={
prop: val[bidx]
for prop, val in material["texture_maps"].items()
},
o3d_type=o3d_type,
connection=buf_con):
raise IOError(
"[Open3D set_mesh_data] Geometry data serialization for tag "
"{tag} step {step} failed!")
# TODO(ssheorey): This returns a copy instead of the original. Benchmark
# vs numpy
data_buffer = buf_con.get_buffer()
filename, this_write_location = _async_data_writer.enqueue(
os.path.join(write_dir, tag.replace('/', '-')), data_buffer)
if bidx == 0:
geometry_metadata.batch_index.filename = filename
geometry_metadata.batch_index.start_size.add(
start=this_write_location,
size=len(data_buffer),
masked_crc32c=masked_crc32c(data_buffer))
if data_bbox is not None:
data_bbox_proto = plugin_data_pb2.InferenceData()
for l, c in zip(data_bbox['bbox_labels'][bidx],
data_bbox['bbox_confidences'][bidx]):
data_bbox_proto.inference_result.add(label=l, confidence=c)
data_bbox_serial = data_bbox_proto.SerializeToString()
filename, this_write_location = _async_data_writer.enqueue(
os.path.join(write_dir, tag.replace('/', '-')),
data_bbox_serial)
geometry_metadata.batch_index.start_size[
-1].aux_start = this_write_location
geometry_metadata.batch_index.start_size[-1].aux_size = len(
data_bbox_serial)
geometry_metadata.batch_index.start_size[
-1].aux_masked_crc32c = masked_crc32c(data_bbox_serial)
return geometry_metadata.SerializeToString()
def _write_geometry_data(write_dir, tag, step, data, max_outputs=3):
"""Serialize and write geometry data for a tag. Data is written to a
separate file per tag.
TODO: Add version specific reader / writer
Args:
write_dir (str): Path of folder to write data file.
tag (str): Full name for geometry.
step (int): Iteration / step count.
data (dict): Property name to tensor mapping.
max_outputs (int): Only the first `max_samples` data points in each
batch will be saved.
Returns:
A comma separated data location string with the format
f"{filename},{write_location},{write_size}"
"""
if not isinstance(data, dict):
raise TypeError(
"data should be a dict of geometry property names and tensors.")
unknown_props = [
prop for prop in data if prop not in metadata.GEOMETRY_PROPERTY_DIMS
]
if unknown_props:
raise ValueError(
f"Unknown geometry properties in data: {unknown_props}")
if "vertex_positions" not in data:
raise ValueError("Primary key 'vertex_positions' not provided.")
if max_outputs < 1:
raise ValueError(
f"max_outputs ({max_outputs}) should be a non-negative integer.")
max_outputs = int(max_outputs)
batch_size = None
n_vertices = None
n_triangles = None
n_lines = None
vertex_data = {}
triangle_data = {}
line_data = {}
geometry_metadata = plugin_data_pb2.Open3DPluginData(
version=metadata._VERSION)
o3d_type = "PointCloud"
for prop, tensor in data.items():
if prop in ('vertex_positions', ) + metadata.VERTEX_PROPERTIES:
prop_name = prop[7:]
vertex_data[prop_name] = _preprocess(prop, tensor, step,
max_outputs,
geometry_metadata)
if vertex_data[prop_name] is None: # Step reference
del vertex_data[prop_name]
continue
if batch_size is None: # Get tensor dims from earlier property
batch_size, n_vertices, _ = vertex_data[prop_name].shape
exp_shape = (batch_size, n_vertices,
metadata.GEOMETRY_PROPERTY_DIMS[prop])
if tuple(vertex_data[prop_name].shape) != exp_shape:
raise ValueError(
f"Property {prop} tensor should be of shape "
f"{exp_shape} but is {vertex_data[prop_name].shape}.")
elif prop in ('triangle_indices', ) + metadata.TRIANGLE_PROPERTIES:
o3d_type = "TriangleMesh"
prop_name = prop[9:]
triangle_data[prop_name] = _preprocess(prop, tensor, step,
max_outputs,
geometry_metadata)
if triangle_data[prop_name] is None: # Step reference
del triangle_data[prop_name]
continue
if n_triangles is None: # Get tensor dims from earlier property
_, n_triangles, _ = triangle_data[prop_name].shape
exp_shape = (batch_size, n_triangles,
metadata.GEOMETRY_PROPERTY_DIMS[prop])
if tuple(triangle_data[prop_name].shape) != exp_shape:
raise ValueError(
f"Property {prop} tensor should be of shape "
f"{exp_shape} but is {triangle_data[prop_name].shape}.")
elif prop in ('line_indices', ) + metadata.LINE_PROPERTIES:
if o3d_type != "TriangleMesh":
o3d_type = "LineSet"
prop_name = prop[5:]
line_data[prop_name] = _preprocess(prop, tensor, step, max_outputs,
geometry_metadata)
if line_data[prop_name] is None: # Step reference
del line_data[prop_name]
continue
if n_lines is None: # Get tensor dims from earlier property
_, n_lines, _ = line_data[prop_name].shape
exp_shape = (batch_size, n_lines,
metadata.GEOMETRY_PROPERTY_DIMS[prop])
if tuple(line_data[prop_name].shape) != exp_shape:
raise ValueError(
f"Property {prop} tensor should be of shape "
f"{exp_shape} but is {line_data[prop_name].shape}.")
vertices = vertex_data.pop("positions",
o3d.core.Tensor((), dtype=o3d.core.float32))
faces = triangle_data.pop("indices",
o3d.core.Tensor((), dtype=o3d.core.int32))
lines = line_data.pop("indices", o3d.core.Tensor((), dtype=o3d.core.int32))
for bidx in range(batch_size):
buf_con = o3d.io.rpc.BufferConnection()
if not o3d.io.rpc.set_mesh_data(
path=tag,
time=step,
layer="",
vertices=vertices[bidx, :, :]
if vertices.ndim == 3 else vertices,
vertex_attributes={
prop: tensor[bidx, :, :]
for prop, tensor in vertex_data.items()
},
faces=faces[bidx, :, :] if faces.ndim == 3 else faces,
face_attributes={
prop: tensor[bidx, :, :]
for prop, tensor in triangle_data.items()
},
lines=lines[bidx, :, :] if lines.ndim == 3 else lines,
line_attributes={
prop: tensor[bidx, :, :]
for prop, tensor in line_data.items()
},
o3d_type=o3d_type,
connection=buf_con):
raise IOError(
"[Open3D set_mesh_data] Geometry data serialization for tag "
"{tag} step {step} failed!")
# TODO(ssheorey): This returns a copy instead of the original. Benchmark
# vs numpy
data_buffer = buf_con.get_buffer()
filename, this_write_location = _async_data_writer.enqueue(
os.path.join(write_dir, tag.replace('/', '-')), data_buffer)
if bidx == 0:
geometry_metadata.batch_index.filename = filename
geometry_metadata.batch_index.start_size.add(
start=this_write_location,
size=len(data_buffer),
masked_crc32c=masked_crc32c(data_buffer))
return geometry_metadata.SerializeToString()
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