isdigit

2 files changed, 6 insertions, 2 deletions

diff --git a/augment.py b/augment.py
index 4aab5f1..87d5717 100755
--- a/augment.py
+++ b/augment.py
@@ -80,7 +80,11 @@ if opt.which_epoch == 'latest':
 else:
   current_epoch = opt.which_epoch
 
-epoch_id = "{:02d}_{:04d}_{:04d}".format(current_epoch, data_opt.augment_take, data_opt.augment_make)
+if str(current_epoch).isdigit():
+  epoch_id = "{:02d}_{:04d}_{:04d}".format(int(current_epoch), data_opt.augment_take, data_opt.augment_make)
+else:
+  epoch_id = "{}_{:04d}_{:04d}".format(current_epoch, data_opt.augment_take, data_opt.augment_make)
+
 opt.render_dir = os.path.join(opt.results_dir, opt.name, epoch_id)
 
 if not opt.engine and not opt.onnx:
diff --git a/build_dataset.pl b/build_dataset.pl
index f88b818..9d77766 100755
--- a/build_dataset.pl
+++ b/build_dataset.pl
@@ -39,7 +39,7 @@ my $dir;
 my $src;
 my $dst;
 for ($i = 0; $i < $count; $i++) {
-  if (($i % $test_split) = ($test_split-1)) {
+  if (($i % $test_split) == ($test_split-1)) {
     $dir = $thumbs_dir . "test_";
 
     $src = $images_dir . $images[$i];