update test.py --save-txt

fhkiel-mlaip · Jul 16, 2020 · 9de5a7a · 9de5a7a
1 parent 825e729
commit 9de5a7a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/test.py b/test.py
@@ -27,7 +27,12 @@ def test(data,
 
     else:  # called directly
         device = torch_utils.select_device(opt.device, batch_size=batch_size)
-        merge = opt.merge  # use Merge NMS
+        merge, save_txt = opt.merge, opt.save_txt  # use Merge NMS, save *.txt labels
+        if save_txt:
+            out = Path('inference/output')
+            if os.path.exists(out):
+                shutil.rmtree(out)  # delete output folder
+            os.makedirs(out)  # make new output folder
 
         # Remove previous
         for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')):