feat: better handling of custom mitsuba3 installation

diegoroyo · Jul 31, 2024 · 0eb0802 · 0eb0802
1 parent de52584
commit 0eb0802
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 16 deletions.
diff --git a/tal/__main__.py b/tal/__main__.py
@@ -140,16 +140,28 @@ def main():
             ask_for_config(
                 Config.MITSUBA2_TRANSIENT_NLOS_FOLDER, force_ask=True)
         elif version == '3':
-            log(LogLevel.INFO, 'Checking if mitsuba can be imported...')
-            try:
-                import mitsuba
-                log(LogLevel.INFO, 'OK mitsuba can be imported')
-            except ModuleNotFoundError:
-                log(LogLevel.PROMPT, 'mitsuba cannot be imported. '
-                    'This can happen if you run a custom installation. '
-                    'Please point to the mitsuba3 (not mitransient) folder:')
-                ask_for_config(
-                    Config.MITSUBA3_TRANSIENT_NLOS_FOLDER, force_ask=True)
+            ask_for_config(
+                Config.MITSUBA3_TRANSIENT_NLOS_FOLDER, force_ask=True)
+            found = False
+            while not found:
+                log(LogLevel.INFO, 'Checking if Mitsuba 3 can be found...')
+                try:
+                    import mitsuba as mi
+                    log(LogLevel.INFO, f'OK Mitsuba 3 can be found (without setpath.sh at {mi.__file__})')
+                    found = True
+                except ModuleNotFoundError:
+                    from tal.render.mitsuba3_transient_nlos import add_mitsuba_to_path
+                    add_mitsuba_to_path()
+                    try:
+                        import mitsuba as mi
+                        log(LogLevel.INFO, f'OK Mitsuba 3 can be found (with setpath.sh at {mi.__file__})')
+                        found = True
+                    except ModuleNotFoundError:
+                        log(LogLevel.PROMPT, 'Mitsuba 3 cannot be found. '
+                            'Please install Mitsuba 3 using \'pip install mitsuba\' (and write \'pip\' in the next prompt) '
+                            'or compile Mitsuba 3 and point to your installation folder:')
+                        ask_for_config(
+                            Config.MITSUBA3_TRANSIENT_NLOS_FOLDER, force_ask=True)
         else:
             raise AssertionError(
                 f'Invalid MITSUBA_VERSION={version}, must be one of (2, 3)')

diff --git a/tal/config.py b/tal/config.py
@@ -26,9 +26,9 @@ class Config(Enum):
          lambda s: os.path.isdir(s))
     MITSUBA3_TRANSIENT_NLOS_FOLDER = \
         ('MITSUBA3_TRANSIENT_NLOS_FOLDER',
-         'Location of Mitsuba 3 (not mitransient) installation folder',
+         'Path to your compiled Mitsuba 3 (not mitransient). If you installed mitsuba from pip, write \'pip\'',
          '',
-         lambda s: os.path.isdir(s))
+         lambda s: os.path.isdir(s) or s == 'pip')
     LOG_LEVEL = \
         ('LOG_LEVEL',
          f'Logging level ({", ".join(LogLevel.__members__)})',
@@ -75,6 +75,10 @@ def ask_for_config(param_name: Config, force_ask=True):
                 param_value = default_value
             if not is_valid(param_value):
                 log(LogLevel.PROMPT, 'Invalid value.')
+        # custom interaction for MITSUBA3_TRANSIENT_NLOS parameter
+        # where the user notifies that it has installed mitsuba 3 through pip
+        if param_value == 'pip':
+            param_value = ''
         config_dict[param_name] = param_value
         write_config(config_dict)
     return config_dict[param_name]

diff --git a/tal/render/mitsuba3_transient_nlos.py b/tal/render/mitsuba3_transient_nlos.py
@@ -22,9 +22,7 @@ def _get_setpath_location():
     return setpath_location
 
 
-try:
-    import mitsuba
-except ModuleNotFoundError:
+def add_mitsuba_to_path():
     import sys
     import subprocess
     setpath_location = _get_setpath_location()
@@ -36,7 +34,13 @@ def _get_setpath_location():
         (key, _, value) = line.partition('=')
         if key == 'PYTHONPATH':
             for directory in value.split(':')[:-1]:
-                sys.path.append(directory)
+                sys.path.insert(0, directory)
+
+
+from tal.config import read_config, Config
+custom_path = read_config().get(Config.MITSUBA3_TRANSIENT_NLOS_FOLDER.value[0], None)
+if custom_path:
+    add_mitsuba_to_path()
 
 
 def get_name():