diff --git a/README.md b/README.md index 190c594..c74c475 100644 --- a/README.md +++ b/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 diff --git a/run.py b/run.py index b080dc0..cc055f1 100644 --- a/run.py +++ b/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: