Merge pull request #5 from corkborg/fix-overwrite

Support overwrite option.

README.md

@@ -13,9 +13,7 @@ pip install -r requirements.txt
 ## usage
 
 ```
-usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT]
-              [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
-              [--overwrite OVERWRITE]
+usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
 
 options:
   -h, --help            show this help message and exit
@@ -24,10 +22,9 @@ options:
   --threshold THRESHOLD
                         Prediction threshold (default is 0.35)
   --ext EXT             Extension to add to caption file in case of dir option (default is .txt)
+  --overwrite           Overwrite the file if it exists.
   --model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}
                         modelname to use for prediction (default is wd14-convnextv2.v1)
-  --overwrite OVERWRITE
-                        Overwrite the file if it exists (default is True)
 ```
 
 single file

run.py

@@ -23,16 +23,15 @@ parser.add_argument(
     '--ext',
     default='.txt',
     help='Extension to add to caption file in case of dir option (default is .txt)')
+parser.add_argument(
+    '--overwrite',
+    action='store_true',
+    help='Overwrite the file if it exists')
 parser.add_argument(
     '--model',
     default='wd14-convnextv2.v1',
     choices=list(interrogators.keys()),
     help='modelname to use for prediction (default is wd14-convnextv2.v1)')
-parser.add_argument(
-    '--overwrite',
-    default=True,
-    type=bool,
-    help='Overwrite the file if it exists (default is True)')
 args = parser.parse_args()
 
 # get interrogator configs
@@ -51,11 +50,20 @@ if args.dir:
     for f in d.iterdir():
         if not f.is_file() or f.suffix not in ['.png', '.jpg', '.jpeg', '.webp']:
             continue
+
         image_path = Path(f)
+        caption_path = image_path.parent / f"{f.stem}{args.ext}"
+
+        if  caption_path.is_file() and not args.overwrite:
+            # skip if file exists
+            print('skip:', image_path)
+            continue
+
         print('processing:', image_path)
         tags = image_interrogate(image_path)
         tags_str = ", ".join(tags.keys())
-        with open(f.parent / f"{f.stem}{args.ext}", "w") as fp:
+
+        with open(caption_path, "w") as fp:
             fp.write(tags_str)
 
 if args.file: