Update graphic.py

Add log_to_filename
pedrodiamel · Aug 21, 2023 · c59034e · c59034e
1 parent 7897dff
commit c59034e
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/pytvision/graphic.py b/pytvision/graphic.py
@@ -6,21 +6,21 @@
 
 
 class VisdomLogger(object):
-    def __init__(self, env_name="main", server=None, port=None) -> None:
+    def __init__(self, env_name="main", log_to_filename=None, server=None, port=None) -> None:
         if server is None:
             server = os.environ.get("VISDOM_SERVER_URL", "localhost")
         if port is None:
             port = int(os.environ.get("VISDOM_PORT", 6006))
 
-        self.viz = Visdom(server=server, port=port, use_incoming_socket=False)
+        self.viz = Visdom(log_to_filename=log_to_filename, server=server, port=port, use_incoming_socket=False)
         self.env = env_name
 
 
 class VisdomLinePlotter(VisdomLogger):
     """Plots to Visdom"""
 
-    def __init__(self, env_name="main", server=None, port=None):
-        super(VisdomLinePlotter, self).__init__(env_name, server, port)
+    def __init__(self, env_name="main", log_to_filename=None, server=None, port=None):
+        super(VisdomLinePlotter, self).__init__(env_name, log_to_filename, server, port)
         self.plots = {}
 
     def plot(self, var_name, split_name, x, y):
@@ -59,8 +59,8 @@ def plot(self, var_name, split_name, x, y):
 class VisdomScatter(VisdomLogger):
     """Scatter to Visdom"""
 
-    def __init__(self, env_name="main", server=None, port=None):
-        super(VisdomScatter, self).__init__(env_name, server, port)
+    def __init__(self, env_name="main", log_to_filename=None, server=None, port=None):
+        super(VisdomScatter, self).__init__(env_name, log_to_filename, server, port)
         self.scatters = {}
 
     def scatter(self, X, Y, title, legend, markersize=10):
@@ -92,8 +92,8 @@ def scatter(self, X, Y, title, legend, markersize=10):
 class HeatMapVisdom(VisdomLogger):
     """Heat Map to Visdom"""
 
-    def __init__(self, env_name="main", heatsize=None, server=None, port=None):
-        super(HeatMapVisdom, self).__init__(env_name, server, port)
+    def __init__(self, env_name="main", heatsize=None, log_to_filename=None, server=None, port=None):
+        super(HeatMapVisdom, self).__init__(env_name, log_to_filename, server, port)
         self.hmaps = {}
         self.heatsize = heatsize
 
@@ -110,8 +110,8 @@ def show(self, title, image):
 class ImageVisdom(VisdomLogger):
     """Images to Visdom"""
 
-    def __init__(self, env_name="main", imsize=None, server=None, port=None):
-        super(ImageVisdom, self).__init__(env_name, server, port)
+    def __init__(self, env_name="main", imsize=None, log_to_filename=None, server=None, port=None):
+        super(ImageVisdom, self).__init__(env_name, log_to_filename, server, port)
         self.images = {}
         self.imsize = imsize