update test.py --save-txt

5 years ago · 37acbdc0b6
parent b8c2da4a0d
commit 37acbdc0b6
1 changed files with 6 additions and 1 deletions
--- 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')):