# ---------------------------------------------------------------------------- # - Open3D: www.open3d.org - # ---------------------------------------------------------------------------- # Copyright (c) 2018-2023 www.open3d.org # SPDX-License-Identifier: MIT # ---------------------------------------------------------------------------- """Utility functions for the Open3D TensorBoard plugin.""" import os from glob import iglob from copy import deepcopy from collections import OrderedDict import logging import threading import numpy as np from tensorboard.backend.event_processing.plugin_event_multiplexer import EventMultiplexer from tensorboard.backend.event_processing.plugin_asset_util import PluginDirectory from tensorboard.compat.tensorflow_stub.pywrap_tensorflow import masked_crc32c import open3d as o3d from open3d.visualization import rendering # TODO(@ssheorey) Colormap and LabelLUT are duplicated from Open3D-ML. Remove # duplicates when 3DML is available on Windows. from .colormap import Colormap from .labellut import LabelLUT from . import plugin_data_pb2 from . import metadata try: from tensorflow.io.gfile import GFile as _fileopen except ImportError: _fileopen = open # Setup Python logger to emulate Open3D C++ logger. _log = logging.getLogger("Open3D") _log.propagate = False _stream_handler = logging.StreamHandler() _stream_handler.setFormatter( logging.Formatter('[%(name)s %(levelname)s T:%(threadName)s] %(message)s')) _stream_handler.setLevel(logging.WARNING) _log.setLevel(logging.WARNING) _log.addHandler(_stream_handler) class ReadWriteLock: """A lock object that allows many simultaneous "read locks", but only one "write lock." Implementation from Python Cookbook (O'Reilly) https://www.oreilly.com/library/view/python-cookbook/0596001673/ch06s04.html Credit: Sami Hangaslammi """ def __init__(self): self._read_ready = threading.Condition(threading.Lock()) self._readers = 0 def acquire_read(self): """Acquire a read lock. Blocks only if a thread has acquired the write lock. """ self._read_ready.acquire() try: self._readers += 1 finally: self._read_ready.release() def release_read(self): """Release a read lock.""" self._read_ready.acquire() try: self._readers -= 1 if not self._readers: self._read_ready.notify_all() finally: self._read_ready.release() def acquire_write(self): """Acquire a write lock. Blocks until there are no acquired read or write locks. """ self._read_ready.acquire() while self._readers > 0: self._read_ready.wait() def release_write(self): """Release a write lock.""" self._read_ready.release() class LRUCache: """Cache storing recent items, i.e. least recently used will be ejected when a new item needs to be added to a full cache. This is thread safe for concurrent access. """ def __init__(self, max_items=128): """ Args: max_items (int): Max items in cache. """ self.cache = OrderedDict() self.max_items = max_items self.rwlock = ReadWriteLock() # hits, misses are not protected against concurrent access for # performance. self.hits = 0 self.misses = 0 def get(self, key): """Retrieve value corresponding to ``key`` from the cache. Return: Value if ``key`` is found, else None. """ self.rwlock.acquire_read() value = self.cache.get(key) # None if not found self.rwlock.release_read() if value is None: self.misses += 1 _log.debug(str(self)) return None self.rwlock.acquire_write() self.cache.move_to_end(key) self.rwlock.release_write() self.hits += 1 _log.debug(str(self)) return value def put(self, key, value): """Add (key, value) pair to the cache. If cache limits are exceeded, eject key-value pairs till the cache is within limits. """ self.rwlock.acquire_write() self.cache[key] = value self.cache.move_to_end(key) if len(self.cache) > self.max_items: self.cache.popitem(last=False) self.rwlock.release_write() _log.debug(str(self)) def clear(self): """Invalidate cache.""" self.rwlock.acquire_write() self.cache.clear() self.rwlock.release_write() def __str__(self): return (f"Items: {len(self.cache)}/{self.max_items}, " f"Hits: {self.hits}, Misses: {self.misses}") def _classify_properties(tensormap): """Classify custom geometry properties (other than positions, indices, colors, normals) into labels and other custom properties. Args: tensormap: geometry TensorMap (e.g.: `point` or `vertex`) Returns: label_prop (List): properties treated as labels (1D Int values). custom_prop (Dict): properties treated as custom properties (Float values, possibly multidimensional). """ label_prop = {} custom_prop = {} exp_size = (tensormap.positions.shape[0] if 'positions' in tensormap else tensormap.indices.shape[0]) for name, tensor in tensormap.items(): if name in ('positions', 'colors', 'normals', 'indices') or name.startswith('_'): continue if (tensor.shape == (exp_size, 1) and tensor.dtype not in (o3d.core.float32, o3d.core.float64, o3d.core.undefined)): label_prop.update({name: 1}) else: custom_prop.update({name: tensor.shape[1]}) return label_prop, custom_prop class Open3DPluginDataReader: """Manage TB event data and geometry data for common use by all Open3DPluginWindow instances. This is thread safe for simultaneous use by multiple browser clients with a multi-threaded web server. Read geometry data is cached in memory. Args: logdir (str): TensorBoard logs directory. cache_max_items (int): Max geometry elements to be cached in memory. """ def __init__(self, logdir, cache_max_items=128): self.logdir = logdir self.event_mux = EventMultiplexer(tensor_size_guidance={ metadata.PLUGIN_NAME: 0 # Store all metadata in RAM }) self._run_to_tags = {} self._event_lock = threading.Lock() # Protect TB event file data # Geometry data reading self._tensor_events = dict() self.geometry_cache = LRUCache(max_items=cache_max_items) self.runtag_prop_shape = dict() self._file_handles = {} # {filename, (open_handle, read_lock)} self._file_handles_lock = threading.Lock() self.reload_events() def _have_data(self): """Do we have any Open3D data?""" try: next( iglob(os.path.join(self.logdir, "**", "plugins", "Open3D", "*.msgpack"), recursive=True)) return True except StopIteration: return False def reload_events(self): """Reload event file""" # Quickly check for Open3D data, since adding all runs can take a long # time if not self._have_data(): with self._event_lock: self._run_to_tags = {} _log.debug(f"No event data found in {self.logdir}") else: self.event_mux.AddRunsFromDirectory(self.logdir) self.event_mux.Reload() run_tags = self.event_mux.PluginRunToTagToContent( metadata.PLUGIN_NAME) with self._event_lock: self._run_to_tags = { run: sorted(tagdict.keys()) for run, tagdict in sorted(run_tags.items()) } _log.debug(f"Event data reloaded: {self._run_to_tags}") self._tensor_events = dict() # Invalidate index # Close all open files with self._file_handles_lock: while len(self._file_handles) > 0: unused_filename, file_handle = self._file_handles.popitem() with file_handle[1]: file_handle[0].close() def is_active(self): """Do we have any Open3D data to display?""" with self._event_lock: return self._have_data() and any( len(tags) > 0 for tags in self._run_to_tags.values()) @property def run_to_tags(self): """Locked access to the run_to_tags map.""" with self._event_lock: return self._run_to_tags def tensor_events(self, run): """Locked access to tensor events for a run.""" with self._event_lock: if run not in self._tensor_events: self._tensor_events[run] = { tag: self.event_mux.Tensors(run, tag) for tag in self._run_to_tags[run] } return self._tensor_events[run] def get_label_to_names(self, run, tag): """Get label (id) to name (category) mapping for a tag.""" md_proto = self.event_mux.SummaryMetadata(run, tag) lab2name = metadata.parse_plugin_metadata(md_proto.plugin_data.content) return dict(sorted(lab2name.items())) 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 update_runtag_prop_shape(self, run, tag, geometry, inference_data_proto): """Update list of custom properties and their shapes for different runs and tags. """ tag_prop_shape = self.runtag_prop_shape.setdefault(run, dict()) prop_shape = tag_prop_shape.setdefault(tag, dict()) if len(prop_shape) == 0 and not geometry.is_empty(): for prop_type in ('point', 'vertex'): # exclude 'line' if hasattr(geometry, prop_type): label_props, custom_props = _classify_properties( getattr(geometry, prop_type)) prop_shape.update(custom_props) prop_shape.update(label_props) if len(inference_data_proto.inference_result) > 0: # Only bbox labels can be visualized. Scalars such as # 'confidence' from BoundingBox3D requires # unlitGradient.GRADIENT shader support for LineSet. prop_shape.update({'labels': 1}) 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 _normalize(tensor): """Normalize tensor by scaling and shifting to range [0, 1]. Args: tensor: Open3D / Numpy float tensor. Int tensors are returned unchanged. Return: tuple: (Normalized tensor, min value of tensor, max value of tensor) min and max are stretched to include 0 or 1 for degenerate tensors. """ if tensor.dtype in (o3d.core.float32, o3d.core.float64, np.float32, np.float64): m, M = tensor.min().item(), tensor.max().item() if M <= m + 1e-6: # stretch degenerate range to include [0, 1] m, M = min(m, 0), max(M, 1) return (tensor - m) * (1. / (M - m)), m, M return tensor, 0, 1 def _u8_to_float(color): return tuple(float(c) / 255 for c in color) def _float_to_u8(color): return tuple(round(255 * c) for c in color) class RenderUpdate: """Update rendering for a geometry, based on an incoming "tensorboard/update_rendering" message. Default rendering is used with an empty message. We use uint8 colors everywhere, since they are more compact to store during serialization. Colormaps are converted to float here, since rendering uses float colormaps. """ _LINES_PER_BBOX = 17 # BoundingBox3D ALL_UPDATED = ["property", "shader", "colormap"] _CMAPS = { "RAINBOW": Colormap.make_rainbow(), "GREYSCALE": Colormap.make_greyscale(), } DICT_COLORMAPS = { name: { # float -> uint8, and RGB -> RGBA point.value: _float_to_u8(point.color) + (255,) for point in cmap.points } for name, cmap in _CMAPS.items() } LABELLUT_COLORS = tuple( _float_to_u8(color + [1.0]) for color in LabelLUT.get_colors(mode='lightbg')) def __init__(self, window_scaling, message, label_to_names): from open3d.visualization.async_event_loop import async_event_loop self._gui = async_event_loop self._window_scaling = window_scaling self._label_to_names = label_to_names render_state = message.get("render_state", { "property": "", "index": 0, "shader": "", "colormap": None, }) self._updated = message.get("updated", deepcopy(self.ALL_UPDATED)) self._property = render_state["property"] if int(render_state["index"]) >= 0: self._index = int(render_state["index"]) # For custom scalar / 3-vector visualization self.data_range = None self._shader = render_state["shader"] # { label or value : float tuple (r,g,b,a) in range [0., 1.] } self._colormap = None # Initialize in first apply() if render_state["colormap"] is not None: to_number = int if self._shader == "unlitGradient.LUT" else float self._colormap = { to_number(label_value): tuple(int(c) for c in rgba) for label_value, rgba in sorted(render_state["colormap"].items(), key=lambda l_c: to_number(l_c[0])) } def get_render_state(self): """Return current render state.""" return { "property": self._property, "index": self._index, "shader": self._shader, "colormap": self._colormap, "range": self.data_range or [0, 1] } def _update_range(self, min_val, max_val): if self.data_range is None: self.data_range = [min_val, max_val] else: self.data_range = [ min(min_val, self.data_range[0]), max(max_val, self.data_range[1]) ] def _set_vis_minmax(self, geometry_vertex, material): m, M = 0., 1. if "__visualization_scalar" in geometry_vertex: m = geometry_vertex["__visualization_scalar"].min().item() M = geometry_vertex["__visualization_scalar"].max().item() if M <= m + 1e-6: # stretch degenerate range to include [0, 1] m, M = min(m, 0), max(M, 1) self._update_range(m, M) material.scalar_min, material.scalar_max = m, M _log.debug("material colormap range range set to " f"{material.scalar_min, material.scalar_max}") def _set_render_defaults(self, geometry, inference_result, label_props, custom_props): """Set default options for rendering (shader and property), based on data type and properties. """ geometry_vertex = (geometry.point if hasattr(geometry, 'point') else geometry.vertex) if self._shader == "": if (len(inference_result) > 0 # Have BB labels or len(label_props) > 0): # Have vertex labels self._shader = "unlitGradient.LUT" if self._property == "": self._property = (next(iter(label_props)) if len(label_props) > 0 else "labels") self._index = 0 elif len(custom_props) > 0: self._shader = "unlitGradient.GRADIENT.RAINBOW" if self._property == "": self._property = next(iter(custom_props)) self._index = 0 elif (isinstance(geometry, o3d.t.geometry.LineSet) or "normals" in geometry_vertex): self._shader = "defaultLit" # also proxy for unlitLine elif "colors" in geometry_vertex: self._shader = "defaultUnlit" else: # Only XYZ self._shader = "unlitSolidColor" # Fix incompatible option. TODO(Sameer): Move this logic to JS if (self._property in custom_props and self._shader == "unlitGradient.LUT"): self._shader = "unlitGradient.GRADIENT.RAINBOW" self._colormap = None if (self._property in label_props and self._shader.startswith("unlitGradient.GRADIENT.")): self._shader = "unlitGradient.LUT" self._colormap = None class BackupRestore(): """Manages backup and restore for tensormap[property] tensors. Call backup(tensormap, property) before modifying the property. This will save the data to a backup tensor (with key __property) and ensure a tensor is available for modification / assignment. A list of tensormap and property tuples is maintained so that a single call to restore() will switch the data back to the earlier state. The implementation takes care to only allocate a single backup copy of the tensor data which will be reused on subsequent calls to backup(). """ def __init__(self): self.swap_list = [] # swap these during restore self.backup_list = [] # backup these during restore def backup(self, tm, prop, clone=False, shape=None, dtype=None): """Args: tm (TensorMap): tensormap, such as PointCloud.point prop (str): property key, such as 'indices' clone (bool): Backup should be a clone, else empty tensor. shape, dtype: sape, dtype if property and backup are absent and a new empty tensor should be created. """ if prop in tm: if "__" + prop in tm: # swap tm[prop], tm["__" + prop] = tm["__" + prop], tm[prop] elif clone: tm["__" + prop] = tm[prop].clone() # backup else: tm["__" + prop] = tm[prop] # backup tm[prop] = o3d.core.Tensor.empty(tm["__" + prop].shape, tm["__" + prop].dtype) self.swap_list.append((tm, prop)) return if "__" + prop in tm: # __prop -> prop tm[prop] = tm["__" + prop] del tm["__" + prop] elif shape is not None and dtype is not None: tm[prop] = o3d.core.Tensor.empty(shape, dtype) self.backup_list.append((tm, prop)) def restore(self): show = "tensormap props: (swap_list)" for tm, prop in self.swap_list: # swap tm[prop], tm["__" + prop] = tm["__" + prop], tm[prop] show += prop + repr(tm[prop][:0]) + '\t' + "__" + prop + repr( tm["__" + prop][:0]) self.swap_list = [] show += "\n (backup_list) " for tm, prop in self.backup_list: # backup tm["__" + prop] = tm[prop] del tm[prop] show += "__" + prop + repr(tm["__" + prop][:0]) self.backup_list = [] def apply(self, o3dvis, geometry_name, geometry, inference_data_proto=None): """Apply the RenderUpdate to a geometry. Args: o3dvis (O3DVisualizer): Window containing the geometry. geometry_name (str): Geometry name in the window. geometry (o3d.t.geometry): Geometry whose rendering is to be updated. inference_data_proto : BoundingBox labels and confidences. """ if (len(self._updated) == 0 or geometry.is_empty()): _log.debug("No updates, or empty geometry.") return def get_labelLUT(): """Create {label: color} mapping from list of colors.""" return { label: self.LABELLUT_COLORS[k] for k, label in enumerate(self._label_to_names) } geometry_vertex = (geometry.point if hasattr(geometry, 'point') else geometry.vertex) have_colors = ("colors" in geometry.line if hasattr( geometry, 'line') else ("colors" in geometry.triangle if hasattr( geometry, 'triangle') else "colors" in geometry_vertex)) if inference_data_proto is not None: inference_result = inference_data_proto.inference_result else: inference_result = [] label_props, custom_props = _classify_properties(geometry_vertex) self._set_render_defaults(geometry, inference_result, label_props, custom_props) if o3dvis.scene.has_geometry(geometry_name): updated = self._updated # update_geometry() only accepts tPointCloud if not isinstance(geometry, o3d.t.geometry.PointCloud): o3dvis.remove_geometry(geometry_name) else: updated = deepcopy(self.ALL_UPDATED) geometry_update_flag = 0 material_update_flag = 0 swap = RenderUpdate.BackupRestore() material = o3dvis.get_geometry_material(geometry_name) # Visualize scalar / 3-vector property with color map if "property" in updated or "shader" in updated: # Float Scalar with colormap if ((self._property in custom_props or self._property in label_props) and self._shader.startswith("unlitGradient") and geometry_vertex[self._property].shape[1] > self._index): geometry_vertex["__visualization_scalar"] = geometry_vertex[ self._property][:, self._index].to( o3d.core.float32).contiguous() geometry_update_flag |= rendering.Scene.UPDATE_UV0_FLAG # 3-vector as RGB elif (self._property in custom_props and self._shader == "defaultUnlit" and geometry_vertex[self._property].shape[1] >= 3): swap.backup(geometry_vertex, "colors") geometry_vertex.colors, min_val, max_val = _normalize( geometry_vertex[self._property][:, :3]) self._update_range(min_val, max_val) geometry_update_flag |= rendering.Scene.UPDATE_COLORS_FLAG # Bounding boxes / LineSet if isinstance(geometry, o3d.t.geometry.LineSet): if self._property == "" and self._shader == "unlitSolidColor": swap.backup(geometry.line, "colors", shape=(len(geometry.line.indices), 3), dtype=o3d.core.uint8) geometry.line.colors[:] = next(iter( self._colormap.values()))[:3] elif self._property != "" and self._shader == "unlitGradient.LUT": if self._colormap is None: self._colormap = get_labelLUT() if "colormap" in updated: swap.backup(geometry.line, "indices", clone=True) swap.backup(geometry.line, "colors", shape=(len(geometry.line.indices), 3), dtype=o3d.core.uint8) t_lines = geometry.line.indices t_colors = geometry.line.colors idx = 0 for bbir in inference_result: col = self._colormap.setdefault(bbir.label, (128, 128, 128, 255)) t_colors[idx:idx + self._LINES_PER_BBOX] = col[:3] if col[3] == 0: # alpha t_lines[idx:idx + self._LINES_PER_BBOX] = 0 idx += self._LINES_PER_BBOX # PointCloud, Mesh, LineSet with colors elif (("shader" in updated or "colormap" in updated) and not geometry.has_valid_material()): material_update_flag = 1 material.base_color = ((1.0, 1.0, 1.0, 1.0) if have_colors else (0.5, 0.5, 0.5, 1.0)) if self._shader == "defaultLit": material.shader = "defaultLit" elif self._shader == "defaultUnlit": material.shader = "defaultUnlit" elif self._shader == "unlitSolidColor": material.shader = "unlitSolidColor" if self._colormap is None: self._colormap = {0.0: [128, 128, 128, 255]} material.base_color = _u8_to_float( next(iter(self._colormap.values()))) elif self._shader == "unlitGradient.LUT": # Label Colormap if self._colormap is None: self._colormap = get_labelLUT() material.shader = "unlitGradient" material.gradient = rendering.Gradient() material.gradient.mode = rendering.Gradient.LUT lmin = min(self._label_to_names.keys()) lmax = max(self._label_to_names.keys()) material.scalar_min, material.scalar_max = lmin, lmax if len(self._colormap) > 1: norm_cmap = list((float(label - lmin) / (lmax - lmin), _u8_to_float(color)) for label, color in self._colormap.items()) material.gradient.points = list( rendering.Gradient.Point(*lc) for lc in norm_cmap) else: material.gradient.points = [ rendering.Gradient.Point(0.0, [1.0, 0.0, 1.0, 1.0]) ] # Colormap (RAINBOW / GREYSCALE): continuous data elif self._shader.startswith("unlitGradient.GRADIENT."): if self._colormap is None: self._colormap = deepcopy( self.DICT_COLORMAPS[self._shader[23:]]) material.shader = "unlitGradient" material.gradient = rendering.Gradient() material.gradient.points = list( rendering.Gradient.Point(value, _u8_to_float(color[:3]) + (1.,)) for value, color in self._colormap.items()) material.gradient.mode = rendering.Gradient.GRADIENT self._set_vis_minmax(geometry_vertex, material) if o3dvis.scene.has_geometry(geometry_name): if material_update_flag > 0: self._gui.run_sync( o3dvis.modify_geometry_material, geometry_name, material ) # does not do force_redraw(), so also need update_geometry() self._gui.run_sync(o3dvis.update_geometry, geometry_name, geometry, geometry_update_flag) _log.debug( f"Geometry {geometry_name} updated with flags " f"Geo:{geometry_update_flag:b}, Mat:{material_update_flag:b}") else: self._gui.run_sync(o3dvis.add_geometry, geometry_name, geometry, material if material_update_flag else None) self._gui.run_sync(o3dvis.post_redraw) swap.restore() _log.debug(f"apply complete: {geometry_name} with " f"{self._property}[{self._index}]/{self._shader}") def to_dict_batch(o3d_geometry_list): """Convert sequence of identical Open3D Eigen geometry types to attribute-tensor dictionary. The geometry sequence forms a batch of data. Custom attributes are not supported. Args: o3d_geometry_list (Iterable): Iterable (list / tuple / sequence generator) of Open3D Eigen geometry types. Returns: Dict[str: numpy.array]: Dictionary of property names and corresponding Numpy arrays. """ if len(o3d_geometry_list) == 0: return {} vertex_positions = [] vertex_colors = [] vertex_normals = [] triangle_indices = [] triangle_uvs = [] line_colors = [] line_indices = [] if isinstance(o3d_geometry_list[0], o3d.geometry.PointCloud): for geometry in o3d_geometry_list: vertex_positions.append(np.asarray(geometry.points)) vertex_colors.append(np.asarray(geometry.colors)) vertex_normals.append(np.asarray(geometry.normals)) geo_dict = { 'vertex_positions': np.stack(vertex_positions, axis=0), 'vertex_colors': np.stack(vertex_colors, axis=0), 'vertex_normals': np.stack(vertex_normals, axis=0), } elif isinstance(o3d_geometry_list[0], o3d.geometry.TriangleMesh): for geometry in o3d_geometry_list: vertex_positions.append(np.asarray(geometry.vertices)) vertex_colors.append(np.asarray(geometry.vertex_colors)) vertex_normals.append(np.asarray(geometry.vertex_normals)) triangle_indices.append(np.asarray(geometry.triangles)) if geometry.has_triangle_uvs(): triangle_uvs.append( np.asarray(geometry.triangle_uvs).reshape((-1, 3, 2))) geo_dict = { 'vertex_positions': np.stack(vertex_positions, axis=0), 'vertex_colors': np.stack(vertex_colors, axis=0), 'vertex_normals': np.stack(vertex_normals, axis=0), 'triangle_indices': np.stack(triangle_indices, axis=0), } if len(triangle_uvs) > 0: geo_dict.update(triangle_texture_uvs=np.stack(triangle_uvs, axis=0)) elif isinstance(o3d_geometry_list[0], o3d.geometry.LineSet): for geometry in o3d_geometry_list: vertex_positions.append(np.asarray(geometry.points)) line_colors.append(np.asarray(geometry.colors)) line_indices.append(np.asarray(geometry.lines)) geo_dict = { 'vertex_positions': np.stack(vertex_positions, axis=0), 'line_colors': np.stack(line_colors, axis=0), 'line_indices': np.stack(line_indices, axis=0), } else: raise NotImplementedError( f"Geometry type {type(o3d_geometry_list[0])} is not supported yet.") # remove empty arrays for prop in tuple(geo_dict.keys()): if geo_dict[prop].size == 0: del geo_dict[prop] return geo_dict