import torch
import numpy as np
import os.path as osp
import plotly.graph_objects as go
from src.data import Data, NAG, Cluster
from src.transforms import GridSampling3D, SaveNodeIndex
from src.utils import fast_randperm, to_trimmed
from torch_scatter import scatter_mean
from src.utils.color import *
# TODO: To go further with ipwidgets :
# - https://plotly.com/python/figurewidget-app/
# - https://ipywidgets.readthedocs.io/en/stable/
def visualize_3d(
input,
keys=None,
figsize=1000,
width=None,
height=None,
class_names=None,
class_colors=None,
stuff_classes=None,
num_classes=None,
hide_void_pred=False,
voxel=-1,
max_points=50000,
point_size=3,
centroid_size=None,
error_color=None,
centroids=False,
h_edge=False,
h_edge_attr=False,
h_edge_width=None,
v_edge=False,
v_edge_width=None,
gap=None,
radius=None,
center=None,
select=None,
alpha=0.1,
alpha_super=None,
alpha_stuff=0.2,
point_symbol='circle',
centroid_symbol='circle',
colorscale='Agsunset',
**kwargs):
"""3D data interactive visualization.
:param input: `Data` or `NAG` object
:param keys: `List(str)` or `str`
By default, the following attributes will be parsed in `input`
for visualization {`pos`, `rgb`, `y`, `obj`, `semantic_pred`,
`obj_pred`}. Yet, if `input` contains other attributes that you
want to visualize, these can be passed as `keys`. This only
supports point-wise attributes stored as 1D or 2D tensors.
If the tensor contains only 1 channel, the attribute will be
represented with a grayscale colormap. If the tensor contains
2 or 3 channels, these will be represented as RGB, with
an additional all-1 channel if need be. If the tensor contains
more than 3 channels, a PCA projection to RGB will be shown. In
any case, the attribute values will be rescaled with respect to
their statistics before visualization, meaning that colors may
not compare between two different plots
:param figsize: `int`
Figure dimensions will be `(figsize, figsize/2)` if `width` and
`height` are not specified
:param width: `int`
Figure width
:param height: `int`
Figure height
:param class_names: `List(str)`
Names for point labels found in attributes `y` and
`semantic_pred`
:param class_colors: `List(List(int, int, int))`
Colors palette for point labels found in attributes `y` and
`semantic_pred`
:param stuff_classes: `List(int)`
Semantic labels of the classes considered as `stuff` for
instance and panoptic segmentation. If `y` and `obj` are found
in the point attributes, the stuff annotations will appear
accordingly. Otherwise, stuff instance labeling will appear as
any other object
:param num_classes: `int`
Number of valid classes. By convention, we assume
`y ∈ [0, num_classes-1]` are VALID LABELS, while
`y < 0` AND `y >= num_classes` ARE VOID LABELS
:param hide_void_pred: `bool`
Whether predictions on points labeled as VOID be visualized
:param voxel: `float`
Voxel size to subsample the point cloud to facilitate
visualization
:param max_points: `int`
Maximum number of points displayed to facilitate visualization
:param point_size: `int` or `float`
Size of point markers
:param centroid_size: `int` or `float`
Size of superpoint markers
:param error_color: `List(int, int, int)`
Color used to identify mis-predicted points
:param centroids: `bool`
Whether superpoint centroids should be displayed
:param h_edge: `bool`
Whether horizontal edges should be displayed (only if
`centroids=True`)
:param h_edge_attr: `bool`
Whether the edges should be colored by their features found in
`edge_attr` (only if `h_edge=True`)
:param h_edge_width: `float`
Width of the horizontal edges, if `h_edge=True`. Defaults to
`None`, in which case `point_size` will be used for the edge
width
:param v_edge: `bool`
Whether vertical edges should be displayed (only if
`centroids=True` and `gap` is not `None`)
:param v_edge_width: `float`
Width of the vertical edges, if `v_edge=True`. Defaults to
`None`, in which case `point_size` will be used for the edge
width
:param gap: `List(float, float, float)`
If `None`, the hierarchical graphs will be overlaid on the points.
If not `None`, a 3D tensor indicating the offset by which the
hierarchical graphs should be plotted
:param radius: `float`
If not `None`, only visualize a spherical sampling of the input
data, centered on `center` and with size `radius`. This option
is not compatible with `select`
:param center: `List(float, float, float)`
If `radius` is provided, only visualize a spherical sampling of
the input data, centered on `center` and with size `radius`. If
`None`, the center of the scene will be used
:param select: `Tuple(int, Tensor)`
If not `None`, will call `Data.select(*select)` or
`NAG.select(*select)` on the input data (depending on its nature)
and the coloring schemes will illustrate it. This option is not
compatible with `radius`
:param alpha: `float`
Rules the whitening of selected points, nodes and edges (only if
select is not `None`)
:param alpha_super: `float`
Rules the whitening of superpoints (only if select is not
`None`). If `None`, alpha will be used as fallback
:param alpha_stuff: `float`
Rules the whitening of stuff points (only if the input
points have `obj` and `semantic_pred` attributes, and
`stuff_classes` or `num_classes` is specified). If `None`,
`alpha` will be used as fallback
:param point_symbol: `str`
Marker symbol used for points. Must be one of
`{'circle', 'circle-open', 'square', 'square-open', 'diamond',
'diamond-open', 'cross', 'x'}`. Defaults to `'circle'`
:param centroid_symbol: `str`
Marker symbol used for centroids. Must be one of
`{'circle', 'circle-open', 'square', 'square-open', 'diamond',
'diamond-open', 'cross', 'x'}`. Defaults to `'circle'`
:param colorscale: `str`
Plotly colorscale used for coloring 1D continuous features. See
https://plotly.com/python/builtin-colorscales for options
:param kwargs
:return:
"""
# Data attributes plotted by default if found in the input
_DEFAULT_KEYS = [
'pos',
'rgb',
'y',
'semantic_pred',
'obj',
'obj_pred',
'x',
'super_sampling',
'super_index']
# assert isinstance(input, (Data, NAG))
gap = torch.tensor(gap) if gap is not None else gap
assert gap is None or gap.shape == torch.Size([3])
assert not (radius and (select is not None)), \
"Cannot use both a `radius` and `select` at once"
# We work on copies of the input data, to allow modified in this
# scope
input = input.clone().cpu()
# If the input is a simple Data object, we convert it to a NAG
input = NAG([input]) if isinstance(input, Data) else input
# If the last level of the NAG has super_index, we manually
# construct an additional Data level and append it to the NAG
if input[input.num_levels - 1].is_sub:
data_last = input[input.num_levels - 1]
sub = Cluster(
data_last.super_index, torch.arange(data_last.num_nodes),
dense=True)
obj = data_last.obj.merge(data_last.super_index) \
if data_last.obj else None
pos = scatter_mean(data_last.pos, data_last.super_index, dim=0)
input = NAG(input.to_list() + [Data(pos=pos, sub=sub, obj=obj)])
is_nag = isinstance(input, NAG)
num_levels = input.num_levels if is_nag else 1
# Make sure alpha is in [0, 1]
alpha = max(0, min(alpha, 1))
alpha_super = max(0, min(alpha_super, 1)) if alpha_super else alpha
alpha_stuff = max(0, min(alpha_stuff, 1)) if alpha_stuff else alpha
# If `radius` is provided, we only visualize a spherical selection
# of size `radius` around the `center`
if radius is not None:
# If no `center` provided, pick the middle of the scene
if center is None:
hi = input[0].pos.max(dim=0).values
lo = input[0].pos.min(dim=0).values
center = (hi + lo) / 2
# For Z, we center on the average Z, because the middle
# value may cause empty samplings for outdoor scenes with
# some very high objects and most of the interesting stuff
# happening near the ground
center[2] = input[0].pos[:, 2].mean()
else:
center = torch.as_tensor(center).cpu()
center = center.view(1, -1)
# Create a mask on level-0 (ie points) to be used for indexing
# the NAG structure
mask = torch.where(
torch.linalg.norm(input[0].pos - center, dim=1) < radius)[0]
# Subselect the hierarchical partition based on the level-0 mask
input = input.select(0, mask)
# If `select` is provided, we will call NAG.select on the input data
# and illustrate the selected/discarded pattern in the figure
if select is not None and is_nag:
# Add an ID to the points before applying NAG.select
nag_temp = input.clone()
for i in range(nag_temp.num_levels):
nag_temp._list[i] = SaveNodeIndex()(nag_temp[i])
# Apply the selection
nag_temp = nag_temp.select(*select)
# Indicate, for each node of the hierarchical graph, whether it
# has been selected
for i in range(num_levels):
selected = torch.zeros(input[i].num_nodes, dtype=torch.bool)
selected[nag_temp[i][SaveNodeIndex.DEFAULT_KEY]] = True
input[i].selected = selected
del nag_temp, selected
elif select is not None and not is_nag:
# Add an ID to the points before applying NAG.select
data_temp = SaveNodeIndex()(Data(pos=input.pos.clone()))
# Apply the selection
data_temp = data_temp.select(select)[0]
# Indicate, for each node of the hierarchical graph, whether it
# has been selected
selected = torch.zeros(input.num_nodes, dtype=torch.bool)
selected[data_temp[SaveNodeIndex.DEFAULT_KEY]] = True
input.selected = selected
del data_temp, selected
elif is_nag:
for i in range(num_levels):
input[i].selected = torch.ones(
input[i].num_nodes, dtype=torch.bool)
else:
input.selected = torch.ones(input.num_nodes, dtype=torch.bool)
# Data_0 accounts for the lowest level of hierarchy, the points
# themselves
data_0 = input[0] if is_nag else input
# Subsample to limit the drawing time
# If the level-0 cloud needs to be voxelized or sampled, a NAG
# structure will be affected too. To maintain NAG consistency, we
# only support 'GridSampling3D' with mode='last' and random sampling
# without replacement. To keep track of the sampled points and index
# the NAG accordingly, we use 'SaveNodeIndex'
idx = torch.arange(data_0.num_points)
# If a voxel size is specified, voxelize the level-0. We first
# isolate the 'pos' and the input indices of data_0 and apply
# voxelization on this. We then recover the original grid-sampled
# points indices to be used with Data.select or NAG.select
if voxel > 0:
data_temp = SaveNodeIndex()(Data(pos=data_0.pos.clone()))
data_temp = GridSampling3D(voxel, mode='last')(data_temp)
idx = data_temp[SaveNodeIndex.DEFAULT_KEY]
del data_temp
# If the cloud is too large with respect to required 'max_points',
# sample without replacement
if idx.shape[0] > max_points:
idx = idx[fast_randperm(idx.shape[0])[:max_points]]
# If a sampling is needed, apply it to the input Data or NAG,
# depending on the structure
if idx.shape[0] < data_0.num_points:
input = input.select(0, idx) if is_nag else input.select(idx)[0]
data_0 = input[0] if is_nag else input
# Round to the cm for cleaner hover info
data_0.pos = (data_0.pos * 100).round() / 100
# Class colors initialization
if class_colors is not None and not isinstance(class_colors[0], str):
class_colors = np.asarray(class_colors)
else:
class_colors = None
# Prepare figure
width = width if width and height else figsize
height = height if width and height else int(figsize / 2)
margin = int(0.02 * min(width, height))
layout = go.Layout(
width=width,
height=height,
scene=dict(aspectmode='data', ), # preserve aspect ratio
margin=dict(l=margin, r=margin, b=margin, t=margin),
uirevision=True)
fig = go.Figure(layout=layout)
# To keep track of which trace should be seen under which mode
# (i.e. button), we build trace_modes. This is a list of dictionaries
# indicating, for each trace (list element), which mode (dict key)
# it should appear in and with which attributes (values are dict of
# parameters for plotly figure updates)
trace_modes = []
i_point_trace = 0
i_unselected_point_trace = 1
# Initialize `void_classes`
void_classes = [num_classes] if num_classes else []
# Draw a trace for position-colored 3D point cloud
mini = data_0.pos.min(dim=0).values
maxi = data_0.pos.max(dim=0).values
colors = (data_0.pos - mini) / (maxi - mini + 1e-6)
colors = rgb_to_plotly_rgb(colors)
data_0.pos_colors = colors
fig.add_trace(
go.Scatter3d(
x=data_0.pos[data_0.selected, 0],
y=data_0.pos[data_0.selected, 1],
z=data_0.pos[data_0.selected, 2],
mode='markers',
marker=dict(
symbol=point_symbol,
size=point_size,
color=colors[data_0.selected]),
hoverinfo='x+y+z+text',
hovertext=None,
showlegend=False,
visible=True, ))
trace_modes.append({
'Position RGB': {
'marker.color': colors[data_0.selected], 'hovertext': None}})
fig.add_trace(
go.Scatter3d(
x=data_0.pos[~data_0.selected, 0],
y=data_0.pos[~data_0.selected, 1],
z=data_0.pos[~data_0.selected, 2],
mode='markers',
marker=dict(
symbol=point_symbol,
size=point_size,
color=colors[~data_0.selected],
opacity=alpha),
hoverinfo='x+y+z+text',
hovertext=None,
showlegend=False,
visible=True, ))
trace_modes.append({
'Position RGB': {
'marker.color': colors[~data_0.selected], 'hovertext': None}})
# Draw a trace for RGB 3D point cloud
if data_0.rgb is not None:
colors = data_0.rgb
colors = rgb_to_plotly_rgb(colors)
data_0.rgb_colors = colors
trace_modes[i_point_trace]['RGB'] = {
'marker.color': colors[data_0.selected], 'hovertext': None}
trace_modes[i_unselected_point_trace]['RGB'] = {
'marker.color': colors[~data_0.selected], 'hovertext': None}
# Color the points with ground truth semantic labels. If labels are
# expressed as histograms, keep the most frequent one
if data_0.y is not None:
y = data_0.y
y = y.argmax(1).numpy() if y.dim() == 2 else y.numpy()
colors = class_colors[y] if class_colors is not None \
else int_to_plotly_rgb(torch.LongTensor(y))
data_0.y_colors = colors
if class_names is None:
text = np.array([f"Class {i}" for i in range(y.max() + 1)])
else:
text = np.array([str.title(c) for c in class_names])
text = text[y]
trace_modes[i_point_trace]['Semantic'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Semantic'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
# Color the points with predicted semantic labels. If labels are
# expressed as histograms, keep the most frequent one
if data_0.semantic_pred is not None:
pred = data_0.semantic_pred
pred = pred.argmax(1).numpy() if pred.dim() == 2 else pred.numpy()
# If the ground truth labels are available, we use them to
# identify void points in the predictions
if data_0.y is not None and hide_void_pred:
# Get the target label
y_gt = data_0.y
y_gt = y_gt.argmax(1) if y_gt.dim() == 2 else y_gt
# Create a mask over the points identifying those whose
# ground truth label is void
is_void = np.zeros(y_gt.max() + 1, dtype='bool')
for i in void_classes:
if i < is_void.shape[0]:
is_void[i] = True
is_void = is_void[y_gt]
# Set the predicted label to void if the ground truth is
# void, this avoids visualizing predictions on void
# labels
pred[is_void] = y_gt[is_void]
colors = class_colors[pred] if class_colors is not None \
else int_to_plotly_rgb(torch.LongTensor(pred))
data_0.pred_colors = colors
if class_names is None:
text = np.array([f"Class {i}" for i in range(pred.max() + 1)])
else:
text = np.array([str.title(c) for c in class_names])
text = text[pred]
trace_modes[i_point_trace]['Semantic Pred.'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Semantic Pred.'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
# Color the points with ground truth instance labels. If semantic
# labels and stuff_classes/void_classes also passed, the stuff/void
# annotations will be treated accordingly
if data_0.obj is not None and (class_names is None or data_0.y is None):
obj = data_0.obj if isinstance(data_0.obj, torch.Tensor) \
else data_0.obj.major(num_classes=num_classes)[0]
colors = int_to_plotly_rgb(obj)
data_0.obj_colors = colors
text = np.array([f"Object {o}" for o in obj])
trace_modes[i_point_trace]['Panoptic'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Panoptic'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
elif data_0.obj is not None:
# Colors and text for thing points
obj = data_0.obj.major(num_classes=num_classes)[0]
colors_thing = int_to_plotly_rgb(obj)
text_thing = np.array([f"Object {o}" for o in obj])
# For simplicity, we merge void_classes into the stuff_classes,
# the expected behavior is the same, except that we will ensure
# that the hover text distinguishes between stuff and void
stuff_classes = stuff_classes if stuff_classes is not None else []
stuff_classes = list(set(stuff_classes).union(set(void_classes)))
# Colors and text for stuff points
y = data_0.y
y = y.argmax(1).numpy() if y.dim() == 2 else y.numpy()
colors_stuff = class_colors[y] if class_colors is not None \
else int_to_plotly_rgb(torch.LongTensor(y))
if class_names is None:
text_stuff = np.array([
f"{'Void' if i in void_classes else 'Stuff'} - Class {i}"
for i in range(y.max() + 1)])
else:
text_stuff = np.array([
f"{'Void' if i in void_classes else 'Stuff'} - {str.title(c)}"
for i, c in enumerate(class_names)])
text_stuff = text_stuff[y]
# Apply alpha-whitening on stuff points
colors_stuff = colors_stuff.astype('float')
white = np.full((colors_stuff.shape[0], 3), 255, dtype='float')
colors_stuff = colors_stuff * alpha_stuff + white * (1 - alpha_stuff)
colors_stuff = colors_stuff.astype('int64')
# Compute mask for stuff points
stuff_classes = np.asarray([i for i in stuff_classes if i <= y.max()])
is_stuff = np.zeros(y.max() + 1, dtype='bool')
for i in stuff_classes:
if i < is_stuff.shape[0]:
is_stuff[i] = True
is_stuff = is_stuff[y]
# Merge thing and stuff colors and text
colors = colors_thing
text = text_thing
colors[is_stuff] = colors_stuff[is_stuff]
text[is_stuff] = text_stuff[is_stuff]
data_0.obj_colors = colors
# Create trace modes
trace_modes[i_point_trace]['Panoptic'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Panoptic'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
# Color the points with predicted instance labels. If semantic
# labels and stuff_classes/void_classes also passed, the
# stuff/void predictions will be treated accordingly. This
# expects `data_0.obj_pred` to be an InstanceData object
if getattr(data_0, 'obj_pred', None) is not None and class_names is None:
obj, _, y = data_0.obj_pred.major(num_classes=num_classes)
colors = int_to_plotly_rgb(obj)
data_0.obj_pred_colors = colors
text = np.array([f"Object {o}" for o in obj])
trace_modes[i_point_trace]['Panoptic Pred.'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Panoptic Pred.'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
elif getattr(data_0, 'obj_pred', None) is not None:
# Colors and text for thing points
obj, _, y = data_0.obj_pred.major(num_classes=num_classes)
colors_thing = int_to_plotly_rgb(obj)
text_thing = np.array([f"Object {o}" for o in obj])
# For simplicity, we merge void_classes into the stuff_classes,
# the expected behavior is the same, except that we will ensure
# that the hover text distinguishes between stuff and void
stuff_classes = stuff_classes if stuff_classes is not None else []
stuff_classes = list(set(stuff_classes).union(set(void_classes)))
# If the ground truth labels are available, we use them to
# identify void points in the predictions
if data_0.y is not None and hide_void_pred:
# Get the target label
y_gt = data_0.y
y_gt = y_gt.argmax(1) if y_gt.dim() == 2 else y_gt
# Create a mask over the points identifying those whose
# ground truth label is void
is_void = np.zeros(y_gt.max() + 1, dtype='bool')
for i in void_classes:
if i < is_void.shape[0]:
is_void[i] = True
is_void = is_void[y_gt]
# Set the predicted label to void if the ground truth is
# void, this avoids visualizing predictions on void
# labels
y[is_void] = y_gt[is_void]
# Colors and text for stuff points
colors_stuff = class_colors[y] if class_colors is not None \
else int_to_plotly_rgb(torch.LongTensor(y))
if class_names is None:
text_stuff = np.array([
f"{'Void' if i in void_classes else 'Stuff'} - Class {i}"
for i in range(y.max() + 1)])
else:
text_stuff = np.array([
f"{'Void' if i in void_classes else 'Stuff'} - {str.title(c)}"
for i, c in enumerate(class_names)])
text_stuff = text_stuff[y]
# Apply alpha-whitening on stuff points
colors_stuff = colors_stuff.astype('float')
white = np.full((colors_stuff.shape[0], 3), 255, dtype='float')
colors_stuff = colors_stuff * alpha_stuff + white * (1 - alpha_stuff)
colors_stuff = colors_stuff.astype('int64')
# Compute mask for stuff points
stuff_classes = np.asarray([i for i in stuff_classes if i <= y.max()])
is_stuff = np.zeros(y.max() + 1, dtype='bool')
for i in stuff_classes:
if i < is_stuff.shape[0]:
is_stuff[i] = True
is_stuff = is_stuff[y]
# Merge thing and stuff colors and text
colors = colors_thing
text = text_thing
colors[is_stuff] = colors_stuff[is_stuff]
text[is_stuff] = text_stuff[is_stuff]
data_0.obj_pred_colors = colors
# Create trace modes
trace_modes[i_point_trace]['Panoptic Pred.'] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace]['Panoptic Pred.'] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
# Draw a trace for 3D point cloud features
if data_0.x is not None:
colors = feats_to_plotly_rgb(
data_0.x, normalize=True, colorscale=colorscale)
data_0.x_colors = colors
trace_modes[i_point_trace]['Features 3D'] = {
'marker.color': colors[data_0.selected], 'hovertext': None}
trace_modes[i_unselected_point_trace]['Features 3D'] = {
'marker.color': colors[~data_0.selected], 'hovertext': None}
# Draw a trace for each key specified in keys. Only displays
# point-wise tensor attributes that have not already been plotted
# (ie not in `_DEFAULT_KEYS`)
if keys is None:
keys = []
elif isinstance(keys, str):
keys = [keys]
keys = [k for k in keys if k not in _DEFAULT_KEYS]
for key in keys:
val = getattr(data_0, key, None)
if (val is None or not torch.is_tensor(val)
or val.shape[0] != data_0.num_points):
continue
colors = feats_to_plotly_rgb(
val, normalize=True, colorscale=colorscale)
data_0[f"{key}_colors"] = colors
trace_modes[i_point_trace][str(key).title()] = {
'marker.color': colors[data_0.selected], 'hovertext': None}
trace_modes[i_unselected_point_trace][str(key).title()] = {
'marker.color': colors[~data_0.selected], 'hovertext': None}
# Draw a trace for 3D point cloud sampling (for sampling debugging)
if 'super_sampling' in data_0.keys:
colors = data_0.super_sampling
colors = int_to_plotly_rgb(colors)
colors[data_0.super_sampling == -1] = 230
data_0.super_sampling_colors = colors
trace_modes[i_point_trace]['Super sampling'] = {
'marker.color': colors[data_0.selected], 'hovertext': None}
trace_modes[i_unselected_point_trace]['Super sampling'] = {
'marker.color': colors[~data_0.selected], 'hovertext': None}
# Draw a trace for each cluster level
for i_level, data_i in enumerate(input if is_nag else []):
# Exit in case the Data has no 'super_index'
if not data_i.is_sub:
break
# 'Data.super_index' are expressed between levels i and i+1, but
# we need to recover the 'super_index' between level 0 and i+1,
# to draw clusters on the level-0 points. To this end, we
# compute the desired 'super_index' iteratively, with a
# bottom-up approach
if i_level == 0:
super_index = data_i.super_index
else:
super_index = data_i.super_index[super_index]
# Note that we update the 'trace_modes' 0th element here, this
# assumes only it is the trace holding all level-0 points and on
# which all other colors modes are defined
colors = int_to_plotly_rgb(super_index)
data_0[f"{i_level}_level_colors"] = colors
text = np.array([f"↑: {i}" for i in super_index])
trace_modes[i_point_trace][f"Level {i_level + 1}"] = {
'marker.color': colors[data_0.selected],
'hovertext': text[data_0.selected]}
trace_modes[i_unselected_point_trace][f"Level {i_level + 1}"] = {
'marker.color': colors[~data_0.selected],
'hovertext': text[~data_0.selected]}
# Skip to the next level if we do not need to draw the cluster
# centroids
if not centroids:
continue
# To recover centroids of the i+1 level superpoints, we either
# read them from the next NAG level or compute them using the
# level i 'super_index' indices
num_levels = input.num_levels
is_last_level = i_level == num_levels - 1
if is_last_level or input[i_level + 1].pos is None:
super_pos = scatter_mean(data_0.pos, super_index, dim=0)
else:
super_pos = input[i_level + 1].pos
# Add the gap offset, if need be
if gap is not None:
super_pos += gap * (i_level + 1)
# Round to the cm for cleaner hover info
super_pos = (super_pos * 100).round() / 100
# Save the drawing position of centroids to facilitate vertical
# edges drawing later on
input[i_level + 1].draw_pos = super_pos
# Draw the level-i+1 cluster centroids
idx_sp = torch.arange(data_i.super_index.max() + 1)
colors = int_to_plotly_rgb(idx_sp)
text = np.array([f"#: {i}" for i in idx_sp])
ball_size = centroid_size if centroid_size else point_size * 3
fig.add_trace(
go.Scatter3d(
x=super_pos[input[i_level + 1].selected, 0],
y=super_pos[input[i_level + 1].selected, 1],
z=super_pos[input[i_level + 1].selected, 2],
mode='markers+text',
marker=dict(
symbol=centroid_symbol,
size=ball_size,
color=colors[input[i_level + 1].selected.numpy()],
line_width=min(ball_size / 2, 2),
line_color='black'),
textposition="bottom center",
textfont=dict(size=16),
hovertext=text,
hoverinfo='x+y+z+text',
showlegend=False,
visible=gap is not None, ))
fig.add_trace(
go.Scatter3d(
x=super_pos[~input[i_level + 1].selected, 0],
y=super_pos[~input[i_level + 1].selected, 1],
z=super_pos[~input[i_level + 1].selected, 2],
mode='markers+text',
marker=dict(
symbol=centroid_symbol,
size=ball_size,
color=colors[~input[i_level + 1].selected.numpy()],
line_width=min(ball_size / 2, 2),
line_color='black',
opacity=alpha_super),
textposition="bottom center",
textfont=dict(size=16),
hovertext=text,
hoverinfo='x+y+z+text',
showlegend=False,
visible=gap is not None, ))
keys = [f"Level {i_level + 1}"] if gap is None \
else trace_modes[i_point_trace].keys()
trace_modes.append(
{k: {
'marker.color': colors[input[i_level + 1].selected.numpy()],
'hovertext': text[input[i_level + 1].selected.numpy()]}
for k in keys})
trace_modes.append(
{k: {
'marker.color': colors[~input[i_level + 1].selected.numpy()],
'hovertext': text[~input[i_level + 1].selected.numpy()]}
for k in keys})
if i_level > 0 and v_edge and gap is not None and is_nag:
# Recover the source and target positions for vertical edges
# between i_level -> i_level+1
low_pos = data_i.draw_pos[data_i.selected]
high_pos = super_pos[data_i.super_index[data_i.selected]]
# Convert into a plotly-friendly format for 3D lines
edges = np.full((low_pos.shape[0] * 3, 3), None)
edges[::3] = low_pos
edges[1::3] = high_pos
# Color the vertical edges based on the parent cluster index
# Plotly is a bit hacky with colors for 3D lines. We cannot
# directly pass individual edge colors, we must instead give
# edge color as an int corresponding to a colorscale list
# holding plotly-friendly colors
colors = data_i.super_index[data_i.selected]
colors = np.repeat(colors, 3)
n_colors = colors.max().item() + 1
edge_colorscale = int_to_plotly_rgb(torch.arange(n_colors))
edge_colorscale = [
[i / (n_colors - 1), f"rgb({x[0]}, {x[1]}, {x[2]})"]
for i, x in enumerate(edge_colorscale)]
# Since plotly 3D lines do not support opacity, we draw
# these edges as super thin to limit clutter
edge_width = 0.5 if v_edge_width is None else v_edge_width
# Draw the level i -> i+1 vertical edges. NB we only draw
# edges that are selected and do not draw the unselected
# edges. This is because plotly does not handle opacity
# on lines (yet), which means the unselected edges will tend
# to clutter the figure. For this reason we choose to simply
# not show them
fig.add_trace(
go.Scatter3d(
x=edges[:, 0],
y=edges[:, 1],
z=edges[:, 2],
mode='lines',
line=dict(
width=edge_width,
color=colors,
colorscale=edge_colorscale),
hoverinfo='skip',
showlegend=False,
visible=gap is not None, ))
# NB: at this point, trace_modes contains 'Level i+1' as its
# last key, but we do not want vertical edges to be seen
# when 'Level i+1' is selected, because it means 'Level i'
# nodes are hidden
keys = list(trace_modes[i_point_trace].keys())[:-1]
trace_modes.append({k: {} for k in keys})
# Do not draw superedges if not required or if the i+1 level
# does not have any
if not h_edge or is_last_level or not input[i_level + 1].has_edges:
continue
# Recover the superedge source and target positions
se = input[i_level + 1].edge_index
se_attr = input[i_level + 1].edge_attr
# Since we can only draw one edge direction (they would overlap
# otherwise), we can trim the graph to only keep one direction
# for each undirected edge pair. However, this requires picking
# one direction for the edge attributes to we ARBITRARILY TAKE
# THE MAX EDGE FEATURE for each undirected edge
input[i_level + 1].raise_if_edge_keys()
if se_attr is not None:
se, se_attr = to_trimmed(se, edge_attr=se_attr, reduce='max')
else:
se = to_trimmed(se)
# Recover corresponding source and target coordinates using the
# previously-computed 'super_pos' cluster centroid positions
s_pos = super_pos[se[0]].numpy()
t_pos = super_pos[se[1]].numpy()
# Convert into a plotly-friendly format for 3D lines
edges = np.full((se.shape[1] * 3, 3), None)
edges[::3] = s_pos
edges[1::3] = t_pos
if h_edge_attr and se_attr is not None:
# Recover edge features and convert them to RGB colors. NB:
# edge features are assumed to be in [0, 1] or [-1, 1].
# Since we only draw edges in one direction, we choose to
# only represent the absolute value of the features. This
# implies that features are either direction-independent or
# that the edge direction only changes the sign of the
# feature
colors = feats_to_plotly_rgb(
se_attr.abs(), normalize=True, colorscale=colorscale)
colors = np.repeat(colors, 3, axis=0)
edge_width = point_size if h_edge_width is None else h_edge_width
else:
colors = feats_to_plotly_rgb(
torch.zeros(edges.shape[0]), normalize=True, colorscale=colorscale)
edge_width = point_size if h_edge_width is None else h_edge_width
selected_edge = input[i_level + 1].selected[se].all(axis=0)
selected_edge = selected_edge.repeat_interleave(3).numpy()
# Draw the level-i+1 superedges. NB we only draw edges that are
# selected and do not draw the unselected edges. This is because
# plotly does not handle opacity on lines (yet), which means the
# unselected edges will tend to clutter the figure. For this
# reason we choose to simply not show them
fig.add_trace(
go.Scatter3d(
x=edges[selected_edge, 0],
y=edges[selected_edge, 1],
z=edges[selected_edge, 2],
mode='lines',
line=dict(
width=edge_width,
color=colors[selected_edge]),
hoverinfo='skip',
showlegend=False,
visible=gap is not None, ))
keys = [f"Level {i_level + 1}"] if gap is None \
else trace_modes[i_point_trace].keys()
trace_modes.append({k: {} for k in keys})
# Add a trace for prediction errors. NB: it is important that this
# trace is created last, as the button behavior for this one is
# particular
has_error = data_0.y is not None and data_0.semantic_pred is not None
if has_error:
# Recover prediction and ground truth and deal with potential
# histograms
y = data_0.y
y = y.argmax(1).numpy() if y.dim() == 2 else y.numpy()
pred = data_0.semantic_pred
pred = pred.argmax(1).numpy() if pred.dim() == 2 else pred.numpy()
# Identify erroneous point indices
ignore = void_classes if void_classes else []
ignore = ignore + [-1]
indices = np.where((pred != y) & (~np.in1d(y, ignore)))[0]
# Prepare the color for erroneous points
error_color = 'red' if error_color is None \
else np.asarray[error_color].squeeze()
# Draw the erroneous points
fig.add_trace(
go.Scatter3d(
x=data_0.pos[indices, 0],
y=data_0.pos[indices, 1],
z=data_0.pos[indices, 2],
mode='markers',
marker=dict(
symbol=point_symbol,
size=int(point_size * 1.5),
color=error_color, ),
showlegend=False,
visible=False, ))
# Recover the keys for all visualization modes, as an ordered set,
# with respect to their order of first appearance
modes = list(dict.fromkeys([k for m in trace_modes for k in m.keys()]))
# Traces color for interactive point cloud coloring
def trace_update(mode):
# Prepare the output args for the figure update attributes. By
# default, all traces are non visible, with no color and no
# hover text
n_traces = len(trace_modes)
out = {
'visible': [False] * (n_traces + has_error),
'marker.color': [None] * n_traces,
'hovertext': [''] * n_traces}
# For each trace in 'trace_modes' see if it contains 'mode' and
# adapt out accordingly
for i_trace, t_modes in enumerate(trace_modes):
# The trace has no action for the mode, skip it and leave
# the default args for the trace
if mode not in t_modes:
continue
# Note that a trace will only be visible for its modes
# declared in trace_modes
out['visible'][i_trace] = True
for key, val in t_modes[mode].items():
out[key][i_trace] = val
return [out, list(range(len(trace_modes)))]
# Create the buttons that will serve for toggling trace visibility
updatemenus = [
dict(
buttons=[dict(
label=mode, method='update', args=trace_update(mode))
for mode in modes if mode.lower() != 'errors'],
pad={'r': 10, 't': 10},
showactive=True,
type='dropdown',
direction='right',
xanchor='left',
x=0.02,
yanchor='top',
y=1.02, ),]
if has_error:
updatemenus.append(
dict(
buttons=[dict(
method='restyle',
label='Semantic Errors',
visible=True,
args=[
{'visible': True, 'marker.color': error_color},
[len(trace_modes)]],
args2=[
{'visible': False,},
[len(trace_modes)]],)],
pad={'r': 10, 't': 10},
showactive=False,
type='buttons',
xanchor='left',
x=1.02,
yanchor='top',
y=1.02, ),)
fig.update_layout(updatemenus=updatemenus)
# Place the legend on the left
fig.update_layout(
legend=dict(
yanchor="middle",
y=0.5,
xanchor="right",
x=0.99))
# Hide all axes and no background
fig.update_layout(
scene=dict(
xaxis_title='',
yaxis_title='',
zaxis_title='',
xaxis=dict(
autorange=True,
showgrid=False,
ticks='',
showticklabels=False,
backgroundcolor="rgba(0, 0, 0, 0)"),
yaxis=dict(
autorange=True,
showgrid=False,
ticks='',
showticklabels=False,
backgroundcolor="rgba(0, 0, 0, 0)"),
zaxis=dict(
autorange=True,
showgrid=False,
ticks='',
showticklabels=False,
backgroundcolor="rgba(0, 0, 0, 0)")))
output = {'figure': fig, 'data': data_0}
return output
def figure_html(fig):
# Save plotly figure to temp HTML
fig.write_html(
'/tmp/fig.html',
config={'displayModeBar': False},
include_plotlyjs='cdn',
full_html=False)
# Read the HTML
with open("/tmp/fig.html", "r") as f:
fig_html = f.read()
# Center the figure div for cleaner display
fig_html = fig_html.replace(
'class="plotly-graph-div" style="',
'class="plotly-graph-div" style="margin:0 auto;')
return fig_html
def show(input, path=None, title=None, no_output=True, pt_path=None, **kwargs):
"""Interactive data visualization.
:param input: Data or NAG object
:param path: str
Path to save the visualization into a sharable HTML
:param title: str
Figure title
:param no_output: bool
Set to True if you want to return the 3D Plotly figure objects
:param pt_path:str
Path to save the visualization-ready `Data` object as a `*.pt`.
In this `Data` object, the `pos` and all `*color*` attributes
will be saved, the rest is discarded. This is typically useful
for exporting the visualization layers to another visualization
tool
:param kwargs:
:return:
"""
# Sanitize title and path
if title is None:
title = "Large-scale point cloud"
if path is not None:
if osp.isdir(path):
path = osp.join(path, f"{title}.html")
else:
path = osp.splitext(path)[0] + '.html'
fig_html = f'{title}
'
# Draw a figure for 3D data visualization
out_3d = visualize_3d(input, **kwargs)
if no_output:
if path is None:
out_3d['figure'].show(config={'displayModeBar': False})
else:
fig_html += figure_html(out_3d['figure'])
if path is not None:
with open(path, "w") as f:
f.write(fig_html)
# Save to a .pt file for other downstream tasks.
# NB: we only save the 'pos' and all data attributes containing
# 'color', the rest is discarded
# NB: we save a dictionary, to limit dependencies
if pt_path is not None:
if osp.isdir(pt_path):
pt_path = osp.join(pt_path, f"viz_data.pt")
else:
pt_path = osp.splitext(pt_path)[0] + '.pt'
data = {}
for key in out_3d['data'].keys:
if key == 'pos' or 'color' in key:
data[key] = out_3d['data'][key]
torch.save(data, pt_path)
if not no_output:
return out_3d
return