Kurz vorm durchdrehen

lib/utils/config.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from lib.models.generators.cnn import CNNRouteGeneratorModel, CNNRouteGeneratorDiscriminated
 from lib.models.homotopy_classification.cnn_based import ConvHomDetector
 from lib.utils.model_io import ModelParameters
+from lib.utils.transforms import AsArray
 
 
 def is_jsonable(x):

lib/utils/logging.py

@@ -5,10 +5,12 @@ from pytorch_lightning.loggers.neptune import NeptuneLogger
 from pytorch_lightning.loggers.test_tube import TestTubeLogger
 
 from lib.utils.config import Config
-
+import numpy as np
 
 class Logger(LightningLoggerBase):
 
+    media_dir = 'media'
+
     @property
     def experiment(self):
         if self.debug:
@@ -84,7 +86,9 @@ class Logger(LightningLoggerBase):
 
     def log_image(self, name, image, **kwargs):
         self.neptunelogger.log_image(name, image, **kwargs)
-        image.savefig(self.log_dir / name)
+        step = kwargs.get('step', None)
+        name = f'{step}_{name}' if step is not None else name
+        image.savefig(self.log_dir / self.media_dir / name)
 
     def save(self):
         self.testtubelogger.save()

lib/utils/transforms.py

@@ -8,5 +8,4 @@ class AsArray(object):
 
     def __call__(self, x):
         array = np.zeros((self.width, self.height))
-
         return array