utils/view.py

from copy import deepcopy
from re import S
from typing import Union
import torch
import open3d as o3d
import open3d.visualization.gui as gui
import open3d.visualization.rendering as rendering
from core.opt import MeshOptimizer
import numpy as np
from util.func import to_numpy

from util.snapshot import Snapshot
    

class Viewer:
    def __init__(self, 
            target_vertices:torch.Tensor, #V,3 
            target_faces:torch.Tensor, #F,3 
            snapshots:list[Snapshot],
            vertex_colors:dict[str,list[np.array]]
            ):
        self._target_vertices = target_vertices
        self._target_faces = target_faces
        self._snapshots = snapshots
        self._vertex_colors = vertex_colors

        self._window_o3 = gui.Application.instance.create_window("Continuous Remeshing",1000,800)
        self._window_o3.set_on_layout(self._layout)
        self._scene_widget = gui.SceneWidget()
        self._scene_widget.scene = rendering.Open3DScene(self._window_o3.renderer)
        self._scene_widget.scene.set_background([.5, .5, .5, 1])
        bbox = o3d.geometry.AxisAlignedBoundingBox([-1, -1, -1], [1, 1, 1])
        self._scene_widget.setup_camera(60, bbox, [0, 0, 0])
        self._window_o3.add_child(self._scene_widget)
        self._scene_widget.set_on_mouse(self._on_mouse)

        #lights
        self._scene_widget.scene.scene.enable_sun_light(False)            
        self._scene_widget.scene.scene.set_indirect_light(gui.Application.instance.resource_path + '/park2')
        self._scene_widget.scene.scene.enable_indirect_light(True)
        self._scene_widget.scene.scene.set_indirect_light_intensity(45000)
        self._scene_widget.scene.show_skybox(False)

        #right panel
        margins = gui.Margins(*[self._window_o3.theme.default_margin]*4)
        spacing = self._window_o3.theme.default_layout_spacing
        self._right_panel = gui.Vert(spacing,margins)
        
        def make_checkbox(name,checked):
            checkbox = gui.Checkbox(name)
            checkbox.checked = checked
            checkbox.set_on_checked(lambda *args:self._update())
            self._right_panel.add_child(checkbox)
            return checkbox

        self._mesh_checkbox = make_checkbox("Show Mesh",True)
        self._colorbox = gui.Combobox()
        for item in ['Gray','Relative Velocity nu','Reference Edge Length l_ref',*self._vertex_colors.keys()]:
            self._colorbox.add_item(item)
        self._colorbox.set_on_selection_changed(lambda *args:self._update())
        self._right_panel.add_child(self._colorbox)
        
        self._clim_slider = gui.Slider(gui.Slider.DOUBLE)
        self._clim_slider.double_value = .2
        self._clim_slider.set_limits(1e-3, 1)
        self._clim_slider.set_on_value_changed(lambda *args:self._update())
        self._right_panel.add_child(self._clim_slider)

        self._edges_checkbox = make_checkbox("Show Edges",True)
        self._target_mesh_checkbox = make_checkbox("Show Target Mesh",False)
        self._right_panel.add_child(gui.Label('Ctrl-Click Mesh For Plot!'))
        self._target_edges_checkbox = make_checkbox("Show Target Edges",False)
        self._positions_checkbox = make_checkbox("Plot Positions",False)
        self._gradients_checkbox = make_checkbox("Plot Gradients",False)
        self._m1_checkbox = make_checkbox("Plot m1",False)
        self._m2_checkbox = make_checkbox("Plot m2",False)
        self._nu_checkbox = make_checkbox("Plot nu",True)
        self._lref_checkbox = make_checkbox("Plot l_ref",True)
        self._window_o3.add_child(self._right_panel)

        #bottom panel
        self._bottom_panel = gui.VGrid(cols=2,spacing=spacing)
        self._snapshot_slider = gui.Slider(gui.Slider.INT)
        self._snapshot_slider.int_value = len(self._snapshots)-1
        self._snapshot_slider.set_limits(0, len(self._snapshots)-1)
        self._snapshot_slider.set_on_value_changed(lambda *args:self._update())
        self._bottom_panel.add_child(self._snapshot_slider)
        self._window_o3.add_child(self._bottom_panel)

        self._update()
        
    def _update(self):
        snapshot = self._snapshots[self._snapshot_slider.int_value]
        
        self._scene_widget.scene.clear_geometry()

        self._scene_widget.scene.show_axes(True)
        
        MaterialType = rendering.MaterialRecord if hasattr(rendering,'MaterialRecord') else rendering.Material 

        def add_mesh(name,color,vertices,faces,show_mesh,show_edges,vertex_colors=None):
            vertices_np = vertices.detach().cpu().numpy()
            vertices_o3 = o3d.utility.Vector3dVector(vertices_np)
            faces_o3 = o3d.utility.Vector3iVector(faces.type(torch.int32).cpu().numpy())
            triangleMesh = o3d.geometry.TriangleMesh(vertices_o3,faces_o3)
            triangleMesh.compute_vertex_normals()
            if vertex_colors is not None:
                vertex_colors_np = to_numpy(vertex_colors)
                triangleMesh.vertex_colors = o3d.utility.Vector3dVector(vertex_colors_np)

            if show_mesh:
                material = MaterialType()
                if vertex_colors is None:
                    material.shader = "defaultLit"
                    material.base_color = color
                self._scene_widget.scene.add_geometry(f"{name}_triangleMesh", triangleMesh, material)

            if show_edges:
                edges_material = MaterialType()
                edges_material.base_color = [0,0,0,1]
                edges_material.shader = "unlitLine"
                edges_lineset = o3d.geometry.LineSet.create_from_triangle_mesh(triangleMesh)
                edges_lineset.points = o3d.utility.Vector3dVector(vertices_np + 1e-4 * np.asarray(triangleMesh.vertex_normals))
                self._scene_widget.scene.add_geometry(f"{name}_edges_lineset", edges_lineset, edges_material)

        clim = self._clim_slider.double_value
        if self._colorbox.selected_text=='Relative Velocity nu' and isinstance(snapshot.optimizer, MeshOptimizer):
            vertex_colors = snapshot.optimizer._nu
        elif self._colorbox.selected_text=='Reference Edge Length l_ref' and isinstance(snapshot.optimizer, MeshOptimizer):
            vertex_colors = snapshot.optimizer._ref_len
        elif self._colorbox.selected_text in self._vertex_colors.keys():
            vertex_colors = self._vertex_colors[self._colorbox.selected_text][self._snapshot_slider.int_value]
        else:
            vertex_colors = None

        if vertex_colors is not None:
            c = (to_numpy(vertex_colors) / clim).clip(0,1)
            vertex_colors = np.stack((c,1-c,np.zeros_like(c)),axis=-1)

        add_mesh("mesh",[.5,.5,.5,1],snapshot.vertices,snapshot.faces,self._mesh_checkbox.checked,self._edges_checkbox.checked, vertex_colors)
        add_mesh("target",[.5,.5,1,1],self._target_vertices,self._target_faces,self._target_mesh_checkbox.checked,self._target_edges_checkbox.checked)

    def _layout(self,layout_context):
        r = self._window_o3.content_rect
        
        h = self._bottom_panel.calc_preferred_size(layout_context, gui.Widget.Constraints()).height
        self._bottom_panel.frame = gui.Rect(0, r.height - h, r.width, h)
        r.height -= h

        w = 250
        self._right_panel.frame = gui.Rect(r.width - w, 0, w, r.height)
        r.width -= w

        self._scene_widget.frame = r

    def _on_mouse(self, event):
        if event.type == gui.MouseEvent.Type.BUTTON_DOWN and event.is_modifier_down(gui.KeyModifier.CTRL):
            self._hit_test(event)
            return gui.Widget.EventCallbackResult.HANDLED
        return gui.Widget.EventCallbackResult.IGNORED
        
    def _hit_test(self,event):
        def depth_callback(depth_image):
            f = self._scene_widget.frame
            depth = np.asarray(depth_image)[event.y - f.y, event.x - f.x]
            if depth == 1.0:  # clicked on nothing (i.e. the far plane)
                return
            #need to flip y https://github.com/isl-org/Open3D/issues/4244
            pos = self._scene_widget.scene.camera.unproject(event.x - f.x, f.height - event.y, depth, f.width, f.height)
            
            opt = self._snapshots[self._snapshot_slider.int_value].optimizer
            vertex = (opt.vertices.cpu() - torch.tensor(pos)).norm(dim=-1).argmin().item()
            
            gui.Application.instance.post_to_main_thread(self._window_o3, lambda:self._on_click(pos,vertex))

        self._scene_widget.scene.scene.render_to_depth_image(depth_callback)
        
    def _on_click(self,pos,vertex):
        self._show_plot(vertex)

    def _show_plot(self,vertex):
        ind = self._snapshot_slider.int_value
        device = self._snapshots[0].vertices.device
        vert_ind = torch.zeros(len(self._snapshots),dtype=torch.long,device=device)
        vert_ind[ind] = vertex

        def trace(i):
            nonlocal cur_pos
            vert_ind[i] = (self._snapshots[i].vertices - cur_pos).norm(dim=-1).argmin(dim=0)
            cur_pos = self._snapshots[i].vertices[vert_ind[i]]

        cur_pos = self._snapshots[ind].vertices[vertex] 
        for i in range(ind-1,-1,-1):
            trace(i)

        cur_pos = self._snapshots[ind].vertices[vertex] 
        for i in range(ind+1,len(self._snapshots)):
            trace(i)   

        dims = slice(None,None)

        grad_scale = 100

        from cycler import cycler
        import matplotlib.pyplot as plt
        plt.gca().set_prop_cycle(cycler(linestyle=['-', '--', ':'][dims]))
        
        def extract(prop):
            values = [prop(self._snapshots[i].optimizer,vert_ind[i]) for i in range(0,len(vert_ind))]
            if isinstance(values[0],torch.Tensor):
                values = torch.stack(values).cpu() 
            return values
            
        s = [s.optimizer._step for s in self._snapshots]

        if self._positions_checkbox.checked:
            plt.plot(s,extract(lambda opt,v:opt.vertices[v,dims]),'b',label='pos')
        if self._gradients_checkbox.checked:
            plt.plot(s,grad_scale*extract(lambda opt,v:opt.vertices.grad[v,dims]),'k',label='grad')

        m1 = extract(lambda opt,v:opt._m1[v])
        m2 = extract(lambda opt,v:opt._m2[v])
        velocity = m1 / m2[:,None].sqrt().add_(1e-8) #V,3
        speed = velocity.norm(dim=-1)
        if self._m1_checkbox.checked:
            plt.plot(s,grad_scale*extract(lambda opt,v:opt._m1[v,dims]),'r',label='m1')
        if self._m2_checkbox.checked:
            plt.plot(s,grad_scale*extract(lambda opt,v:opt._m2[v].sqrt()),'-m',label='m2.sqrt()')
        if self._nu_checkbox.checked:
            plt.plot(s,speed,color='orange',label='speed')
            plt.plot(s,extract(lambda opt,v:opt._nu[v]),'-c',label='nu')
        if self._lref_checkbox.checked:
            plt.plot(s,extract(lambda opt,v:opt._ref_len[v]),color='gray',label='l_ref')

        plt.axvline(x=ind, color='k')
        plt.legend()
        plt.grid()
        plt.show()


def show(
    target_vertices:torch.Tensor, #V,3 
    target_faces:torch.Tensor, #F,3 
    snapshots:list[Snapshot],
    vertex_colors:dict[str,list[np.array]]={}
    ):
    for vc in vertex_colors.values():
        assert [c.shape[0] for c in vc] == [s.vertices.shape[0] for s in snapshots]

    gui.Application.instance.initialize()
    viewer = Viewer(target_vertices,target_faces,snapshots,vertex_colors)
    gui.Application.instance.run()