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