[MISC] Faster Genesis import. (#1549)

Author: duburcqa

genesis/engine/mesh.py

@@ -4,10 +4,7 @@
 import fast_simplification
 import numpy as np
 import numpy.typing as npt
-import pyvista as pv
-import tetgen
 import trimesh
-import pymeshlab
 
 import genesis as gs
 from genesis.options.surfaces import Surface
@@ -125,6 +122,9 @@ def remesh(self, edge_len_abs=None, edge_len_ratio=0.01, fix=True):
                 gs.logger.info("Ignoring corrupted cache.")
 
         if not is_cached_loaded:
+            # Importing pymeshlab is very slow and not used very often. Let's delay import.
+            import pymeshlab
+
             gs.logger.info("Remeshing for tetrahedralization...")
             ms = pymeshlab.MeshSet()
             ms.add_mesh(pymeshlab.Mesh(vertex_matrix=self.verts, face_matrix=self.faces))
@@ -151,6 +151,10 @@ def tetrahedralize(self, tet_cfg):
         """
         Tetrahedralize the mesh.
         """
+        # Importing pyvista and tetgen are very slow and not used very often. Let's delay import.
+        import pyvista as pv
+        import tetgen
+
         pv_obj = pv.PolyData(
             self.verts, np.concatenate([np.full((self.faces.shape[0], 1), self.faces.shape[1]), self.faces], axis=1)
         )

genesis/ext/pyrender/renderer.py

@@ -10,7 +10,6 @@
 import pyglet
 import numpy as np
 from OpenGL.GL import *
-import matplotlib.pyplot as plt
 
 from .constants import (
     DEFAULT_Z_FAR,

genesis/ext/pyrender/viewer.py

@@ -16,30 +16,18 @@
 import genesis as gs
 from genesis.vis.rasterizer_context import RasterizerContext
 
+# Importing Tkinter and creating a first context before importing pyglet is necessary to avoid later segfault on MacOS.
+# Note that destroying the window will cause segfault at exit.
 if sys.platform.startswith("darwin"):
-    # Mac OS
     from tkinter import Tk
-    from tkinter import filedialog
-else:
-    try:
-        from Tkinter import Tk
-        from Tkinter import tkFileDialog as filedialog
-    except Exception:
-        try:
-            from tkinter import Tk
-            from tkinter import filedialog as filedialog
-        except Exception:
-            pass
+    from tkinter import filedialog as filedialog
 
-
-try:
     root = Tk()
     root.withdraw()
-except:
-    pass
+else:
+    root = None
 
 import pyglet
-from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter
 from pyglet import clock
 
 from .camera import IntrinsicsCamera, OrthographicCamera, PerspectiveCamera
@@ -917,6 +905,9 @@ def on_key_press(self, symbol: int, modifiers: int) -> EVENT_HANDLE_STATE:
                 self.save_video()
                 self.set_caption(self.viewer_flags["window_title"])
             else:
+                # Importing moviepy is very slow and not used very often. Let's delay import.
+                from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter
+
                 self.video_recorder = FFMPEG_VideoWriter(
                     filename=os.path.join(gs.utils.misc.get_cache_dir(), "tmp_video.mp4"),
                     fps=self.viewer_flags["refresh_rate"],
@@ -1037,6 +1028,8 @@ def _reset_view(self):
         self._trackball = Trackball(self._default_camera_pose, self.viewport_size, scale, centroid)
 
     def _get_save_filename(self, file_exts):
+        global root
+
         file_types = {
             "mp4": ("video files", "*.mp4"),
             "png": ("png files", "*.png"),
@@ -1048,17 +1041,23 @@ def _get_save_filename(self, file_exts):
         save_dir = self.viewer_flags["save_directory"]
         if save_dir is None:
             save_dir = os.getcwd()
+
+        # Importing tkinter is very slow and not used very often. Let's delay import.
+        try:
+            from tkinter import Tk
+            from tkinter import filedialog as filedialog
+        except ImportError:
+            from Tkinter import Tk
+            from Tkinter import tkFileDialog as filedialog
+
         try:
-            master = None
-            if self._run_in_thread:
-                master = Tk()
-                master.withdraw()
+            if root is None:
+                root = Tk()
+                root.withdraw()
             dialog = filedialog.SaveAs(
-                master=master, initialdir=save_dir, title="Select file save location", filetypes=filetypes
+                root, initialdir=save_dir, title="Select file save location", filetypes=filetypes
             )
             filename = dialog.show()
-            if self._run_in_thread:
-                master.destroy()
         except Exception:
             gs.logger.warning("Failed to open file save location dialog.")
             return None

genesis/utils/hybrid.py

@@ -4,11 +4,9 @@
 import time
 from itertools import combinations
 
-import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
 from matplotlib.patches import FancyArrowPatch
-from mpl_toolkits.mplot3d import proj3d
 
 try:
     from pygel3d import graph, hmesh
@@ -220,6 +218,9 @@ def __init__(self, xs, ys, zs, *args, **kwargs):
         self._verts3d = xs, ys, zs
 
     def do_3d_projection(self, renderer=None):
+        # Importing mpl_toolkits is very slow and not used very often. Let's delay import.
+        from mpl_toolkits.mplot3d import proj3d
+
         xs3d, ys3d, zs3d = self._verts3d
         xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
         self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
@@ -238,6 +239,9 @@ def plot_nxgraph(
     node_size=100,
     plot_node_num=True,
 ):
+    # Importing matplotlib is very slow and not used very often. Let's delay import.
+    import matplotlib.pyplot as plt
+
     if pos is None:
         pos = nx.spring_layout(G, dim=3, seed=779)
 

genesis/utils/mesh.py

@@ -14,8 +14,6 @@
 
 import coacd
 import igl
-import pyvista as pv
-import tetgen
 
 import genesis as gs
 
@@ -974,6 +972,10 @@ def make_tetgen_switches(cfg):
 
 
 def tetrahedralize_mesh(mesh, tet_cfg):
+    # Importing pyvista and tetgen are very slow and not used very often. Let's delay import.
+    import pyvista as pv
+    import tetgen
+
     pv_obj = pv.PolyData(
         mesh.vertices, np.concatenate([np.full((mesh.faces.shape[0], 1), mesh.faces.shape[1]), mesh.faces], axis=1)
     )
@@ -1007,6 +1009,9 @@ def visualize_tet(tet, pv_data, show_surface=True, plot_cell_qual=False):
                 scalars=cell_qual, stitle="Quality", cmap="bwr", clim=[0, 1], flip_scalars=True, show_edges=True
             )
         else:
+            # Importing pyvista is very slow and not used very often. Let's delay import.
+            import pyvista as pv
+
             plotter = pv.Plotter()
             plotter.add_mesh(subgrid, "lightgrey", lighting=True, show_edges=True)
             plotter.add_mesh(pv_data, "r", "wireframe")