Fix seed bug and update test code (#46)

synthtiger/gen.py

@@ -33,7 +33,9 @@ def read_config(path):
     return config
 
 
-def generator(path, name, config=None, count=None, worker=0, seed=None, verbose=False):
+def generator(
+    path, name, config=None, count=None, worker=0, seed=None, retry=True, verbose=False
+):
     counter = range(count) if count is not None else itertools.count()
     tasks = _task_generator(seed)
 
@@ -44,7 +46,7 @@ def generator(path, name, config=None, count=None, worker=0, seed=None, verbose=
         post_count = count - pre_count if count is not None else None
 
         for _ in range(worker):
-            _run(_worker, (path, name, config, task_queue, data_queue, verbose))
+            _run(_worker, (path, name, config, task_queue, data_queue, retry, verbose))
         for _ in range(pre_count):
             task_queue.put(next(tasks))
 
@@ -58,7 +60,7 @@ def generator(path, name, config=None, count=None, worker=0, seed=None, verbose=
 
         for _ in counter:
             task_idx, task_seed = next(tasks)
-            data = _generate(template, task_seed, verbose)
+            data = _generate(template, task_seed, retry, verbose)
             yield task_idx, data
 
 
@@ -77,10 +79,10 @@ def set_global_random_states(states):
     imgaug.random.get_global_rng().state = states["imgaug"]
 
 
-def set_global_random_seed(seed):
+def set_global_random_seed(seed=None):
     random.seed(seed)
     np.random.set_state(np.random.RandomState(np.random.MT19937(seed)).get_state())
-    imgaug.seed(seed)
+    imgaug.random.seed(seed)
 
 
 def _run(func, args):
@@ -100,26 +102,28 @@ def _task_generator(seed):
         yield task_idx, task_seed
 
 
-def _worker(path, name, config, task_queue, data_queue, verbose):
+def _worker(path, name, config, task_queue, data_queue, retry, verbose):
     template = read_template(path, name, config)
 
     while True:
         task_idx, task_seed = task_queue.get()
-        data = _generate(template, task_seed, verbose)
+        data = _generate(template, task_seed, retry, verbose)
         data_queue.put((task_idx, data))
 
 
-def _generate(template, seed, verbose):
+def _generate(template, seed, retry, verbose):
     states = get_global_random_states()
     set_global_random_seed(seed)
+    data = None
 
     while True:
         try:
             data = template.generate()
         except:
             if verbose:
                 print(f"{traceback.format_exc()}")
-            continue
+            if retry:
+                continue
         break
 
     set_global_random_states(states)

synthtiger/main.py

@@ -26,6 +26,7 @@ def run(args):
         count=args.count,
         worker=args.worker,
         seed=args.seed,
+        retry=True,
         verbose=args.verbose,
     )
 

tests/conftest.py

@@ -6,34 +6,35 @@
 
 import pytest
 
-import synthtiger
-
 
 @pytest.fixture
-def synthtiger_horizontal_template():
-    """SynthTIGER horizontal template fixture"""
-    config = synthtiger.read_config("examples/synthtiger/config_horizontal.yaml")
-    template = synthtiger.read_template(
-        "examples/synthtiger/template.py", "SynthTiger", config
-    )
-    return template
+def synthtiger_horizontal_args():
+    """SynthTIGER horizontal args"""
+    args = {
+        "script": "examples/synthtiger/template.py",
+        "name": "SynthTiger",
+        "config": "examples/synthtiger/config_horizontal.yaml",
+    }
+    return args
 
 
 @pytest.fixture
-def synthtiger_vertical_template():
-    """SynthTIGER vertical template fixture"""
-    config = synthtiger.read_config("examples/synthtiger/config_vertical.yaml")
-    template = synthtiger.read_template(
-        "examples/synthtiger/template.py", "SynthTiger", config
-    )
-    return template
+def synthtiger_vertical_args():
+    """SynthTIGER vertical args"""
+    args = {
+        "script": "examples/synthtiger/template.py",
+        "name": "SynthTiger",
+        "config": "examples/synthtiger/config_vertical.yaml",
+    }
+    return args
 
 
 @pytest.fixture
-def multiline_template():
-    """Multiline template fixture"""
-    config = synthtiger.read_config("examples/multiline/config.yaml")
-    template = synthtiger.read_template(
-        "examples/multiline/template.py", "Multiline", config
-    )
-    return template
+def multiline_args():
+    """Multiline args"""
+    args = {
+        "script": "examples/multiline/template.py",
+        "name": "Multiline",
+        "config": "examples/multiline/config.yaml",
+    }
+    return args

tests/test_generation_error.py

@@ -4,31 +4,44 @@
 MIT license
 """
 
-import traceback
-
 import pytest
 
+import synthtiger
+
 
 @pytest.mark.parametrize(
-    "template",
+    "args",
     [
-        "synthtiger_horizontal_template",
-        "synthtiger_vertical_template",
-        "multiline_template",
+        "synthtiger_horizontal_args",
+        "synthtiger_vertical_args",
+        "multiline_args",
     ],
 )
-def test_generation_error(template, request):
+def test_generation_error(args, request):
     """Test for errors during data generation"""
 
-    template = request.getfixturevalue(template)
+    args = request.getfixturevalue(args)
+    seed = None
+
+    synthtiger.set_global_random_seed(seed)
+    config = synthtiger.read_config(args["config"])
+    generator = synthtiger.generator(
+        args["script"],
+        args["name"],
+        config=config,
+        count=None,
+        worker=0,
+        seed=seed,
+        retry=False,
+        verbose=True,
+    )
+
     total = 100
     error = 0
 
     for _ in range(total):
-        try:
-            template.generate()
-        except:
+        _, data = next(generator)
+        if data is None:
             error += 1
-            print(traceback.format_exc())
 
     assert error < total