trying to make data generator work from remote

data_generator.py

@@ -4,6 +4,7 @@ import os
 import string
 import subprocess
 import random
+from typing import Iterable
 
 
 class DotDict(dict):
@@ -101,28 +102,32 @@ def generate_equation(latex: DotDict, size, depth=3):
     return equation
 
 
-def generate_image(directory: str, latex_path: str, filename: str, max_length=20):
+def generate_image(directory: str, latex: DotDict, filename: str, max_length=20, equation_depth=3,
+                   pdflatex: str = "/external2/dkkoshman/venv/texlive/2022/bin/x86_64-linux/pdflatex",
+                   ghostscript: str = "/external2/dkkoshman/venv/local/gs/bin/gs"
+                   ):
     """
     Generates a random tex file and corresponding image
     -------
     params:
     :directory: -- dir where to save files
-    :latex_dir: -- path to latex json
-    :filename: -- name for the generated files
+    :latex: -- DotDict with parameters to generate tex
+    :filename: -- absolute filename for the generated files
     :max_length: -- max size of equation
+    :equation_depth: -- max nested level of tex scopes
+    :pdflatex: -- path to pdflatex
+    :ghostscript: -- path to ghostscript
     """
-    # TODO ARGPARSE, path parse
-    filepath = directory + filename
+    # TODO ARGPARSE
+    filepath = os.path.join(directory, filename)
 
-    with open(latex_path) as file:
-        latex = json.load(file)
-    latex = DotDict(latex)
+    equation_length = random.randint(1, max_length)
 
     template = string.Template(latex.template)
     font, font_options = random.choice(latex.fonts)
     font_option = random.choice([''] + font_options)
     fontsize = random.choice(latex.fontsizes)
-    equation = generate_equation(latex, max_length)
+    equation = generate_equation(latex, equation_length, depth=equation_depth)
     tex = template.substitute(font=font, font_option=font_option, fontsize=fontsize, equation=equation)
 
     files_before = set(os.listdir(directory))
@@ -130,7 +135,7 @@ def generate_image(directory: str, latex_path: str, filename: str, max_length=20
         file.write(tex)
 
     pr1 = subprocess.run(
-        f"pdflatex -output-directory={directory} {filepath}.tex".split(),
+        f"{pdflatex} -output-directory={directory} {filepath}.tex".split(),
         stderr=subprocess.PIPE,
     )
 
@@ -138,23 +143,23 @@ def generate_image(directory: str, latex_path: str, filename: str, max_length=20
     if pr1.returncode != 0:
         files_to_delete = files_after - files_before
         if files_to_delete:
-            subprocess.run(['rm'] + [directory + file for file in files_to_delete])
+            subprocess.run(['rm'] + [os.path.join(directory, file) for file in files_to_delete])
         print(pr1.stderr.decode(), tex)
         return
 
     pr2 = subprocess.run(
-        f"gs -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE -o {filepath}.png {filepath}.pdf".split(),
+        f"{ghostscript} -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE -o {filepath}.png {filepath}.pdf".split(),
         stderr=subprocess.PIPE,
     )
 
     files_to_delete = files_after - files_before - {filename + '.png', filename + '.tex'}
     if files_to_delete:
-        subprocess.run(['rm'] + [directory + file for file in files_to_delete])
+        subprocess.run(['rm'] + [os.path.join(directory, file) for file in files_to_delete])
     assert (pr2.returncode == 0)
 
 
 def generate_data(
-        filenames: iter(str),
+        filenames: Iterable[str],
         directory: str,
         latex_path: str,
         overwrite: bool = False
@@ -168,14 +173,27 @@ def generate_data(
     :latex_path: - full path to latex json
     :overwrite: - whether to overwrite existing files
     """
+    subprocess.run(". /external2/dkkoshman/venv/bin/activate")
+    if not os.path.isabs(directory):
+        directory = os.path.join(os.getcwd(), directory)
+    if not os.path.isabs(latex_path):
+        latex_path = os.path.join(os.getcwd(), latex_path)
 
     filenames = set(filenames)
     if not overwrite:
-        existing = set(file.split('.')[0] for file in os.listdir(directory) if file.endswith('.png'))
+        existing = set(
+            filename for file in os.listdir(directory) for filename, ext in os.path.splitext(file) if ext == '.png'
+        )
         filenames -= existing
 
+    with open(latex_path) as file:
+        latex = json.load(file)
+    latex = DotDict(latex)
+
     while filenames:
-        with Pool() as pool:
-            pool.starmap(generate_image, ((directory, latex_path, name) for name in filenames))
+        for name in filenames:
+            generate_image(directory, latex, name)
+        # with Pool() as pool:
+        #     pool.starmap(generate_image, ((directory, latex, name) for name in filenames))
         existing = set(file.split('.')[0] for file in os.listdir(directory) if file.endswith('.png'))
         filenames -= existing

data_preprocessing.py

@@ -23,7 +23,7 @@ class TexImageDataset(Dataset):
         torch.multiprocessing.set_sharing_strategy('file_system')
         self.root_dir = root_dir
         self.filenames = sorted(set(
-            os.path.splitext(filename)[0] for filename in os.listdir(root_dir) if filename.endswith('png')
+            filename for file in os.listdir(root_dir) for filename, ext in os.path.splitext(file) if ext == '.png'
         ))
         self.image_transform = image_transform
         self.tex_transform = tex_transform