Merge pull request #43 from marcpinet/fix-viz2

Fix viz2

neuralnetlib/model.py

@@ -11,7 +11,7 @@
 from neuralnetlib.losses import LossFunction, CategoricalCrossentropy
 from neuralnetlib.optimizers import Optimizer
 from neuralnetlib.preprocessing import PCA
-from neuralnetlib.utils import shuffle, progress_bar, is_interactive
+from neuralnetlib.utils import shuffle, progress_bar, is_interactive, is_display_available
 
 
 class Model:
@@ -126,7 +126,7 @@ def fit(self, x_train: np.ndarray, y_train: np.ndarray, epochs: int, batch_size:
             plot_decision_boundary: Whether to plot the decision boundary
         """
 
-        if plot_decision_boundary and os.getenv('DISPLAY', '') == '' and not is_interactive():
+        if plot_decision_boundary and not is_interactive() and not is_display_available():
             raise ValueError(
                 "Cannot display the plot. Please run the script in an environment with a display.")
 

neuralnetlib/utils.py

@@ -1,6 +1,8 @@
 import sys
 import time
-
+import os
+import platform
+import subprocess
 import numpy as np
 
 
@@ -78,3 +80,33 @@ def is_interactive():
         return not hasattr(main, '__file__')
     except:
         return False
+
+
+def is_display_available():
+    system = platform.system()
+
+    if system == "Linux":
+        return is_display_available_linux()
+    elif system == "Windows":
+        return is_display_available_windows()
+    else:
+        raise NotImplementedError(f"Display check not implemented for {system}")
+
+
+def is_display_available_linux():
+    if "DISPLAY" in os.environ:
+        try:
+            output = subprocess.check_output(["xdpyinfo"], stderr=subprocess.STDOUT)
+            return True
+        except subprocess.CalledProcessError:
+            return False
+    return False
+
+
+def is_display_available_windows():
+    try:
+        import win32api
+        display_devices = win32api.EnumDisplayDevices()
+        return bool(display_devices)
+    except Exception:
+        return False

requirements.txt

@@ -1,2 +1,3 @@
 numpy
-matplotlib  # for plotting the decision boundary (yeah, I won't rewrite matplotlib too)
+matplotlib  # for plotting the decision boundary (yeah, I won't rewrite matplotlib too)
+pywin32; platform_system == 'Windows'  # to check if a display if available (for the plotting)

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='neuralnetlib',
-    version='2.6.1',
+    version='2.6.2',
     author='Marc Pinet',
     description='A simple convolutional neural network library with only numpy as dependency',
     long_description=open('README.md', encoding="utf-8").read(),