diff --git a/README.md b/README.md index 1e9d8c8..3be68f3 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,8 @@ pip install -r requirements.txt ## usage ``` -usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--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}] +usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] + [--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,wd-v1-4-vit-tagger.v3,wd-v1-4-convnext-tagger.v3,wd-v1-4-swinv2-tagger.v3,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] options: -h, --help show this help message and exit @@ -24,7 +25,8 @@ options: --ext EXT Extension to add to caption file in case of dir option (default is .txt) --overwrite Overwrite caption file if it exists --cpu Use CPU only - --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} + --rawtag Use the raw output of the model + --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,wd-v1-4-vit-tagger.v3,wd-v1-4-convnext-tagger.v3,wd-v1-4-swinv2-tagger.v3,mld-caformer.dec-5-97527,mld-tresnetd.6-30000} modelname to use for prediction (default is wd14-convnextv2.v1) ``` @@ -50,6 +52,9 @@ python run.py --file image.jpg --model wd14-convnext.v2 python run.py --file image.jpg --model wd14-convnextv2.v1 python run.py --file image.jpg --model wd14-swinv2-v1 python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2 +python run.py --file image.jpg --model wd-v1-4-vit-tagger.v3 +python run.py --file image.jpg --model wd-v1-4-convnext-tagger.v3 +python run.py --file image.jpg --model wd-v1-4-swinv2-tagger.v3 python run.py --file image.jpg --model mld-caformer.dec-5-97527 python run.py --file image.jpg --model mld-tresnetd.6-30000 ``` diff --git a/run.py b/run.py index a3f45a8..775cf3a 100644 --- a/run.py +++ b/run.py @@ -28,6 +28,10 @@ parser.add_argument( '--cpu', action='store_true', help='Use CPU only') +parser.add_argument( + '--rawtag', + action='store_true', + help='Use the raw output of the model') parser.add_argument( '--model', default='wd14-convnextv2.v1', @@ -47,7 +51,10 @@ def image_interrogate(image_path: Path): """ im = Image.open(image_path) result = interrogator.interrogate(im) - return Interrogator.postprocess_tags(result[1], threshold=args.threshold) + + if args.rawtag: + return Interrogator.postprocess_tags(result[1], threshold=args.threshold) + return Interrogator.postprocess_tags(result[1], threshold=args.threshold, escape_tag=True, replace_underscore=True) if args.dir: d = Path(args.dir) @@ -56,7 +63,7 @@ if args.dir: continue image_path = Path(f) - caption_path = image_path.parent / f"{f.stem}{args.ext}" + caption_path = image_path.parent / f'{f.stem}{args.ext}' if caption_path.is_file() and not args.overwrite: # skip if file exists @@ -65,7 +72,8 @@ if args.dir: print('processing:', image_path) tags = image_interrogate(image_path) - tags_str = ", ".join(tags.keys()) + + tags_str = ', '.join(tags.keys()) with open(caption_path, 'w') as fp: fp.write(tags_str) @@ -73,7 +81,7 @@ if args.dir: if args.file: tags = image_interrogate(Path(args.file)) print() - tags_str = ", ".join(tags.keys()) + tags_str = ', '.join(tags.keys()) print(tags_str) diff --git a/tagger/interrogators.py b/tagger/interrogators.py index d7f1021..cf447be 100644 --- a/tagger/interrogators.py +++ b/tagger/interrogators.py @@ -33,6 +33,18 @@ interrogators: Dict[str, Interrogator] = { 'WD14 moat tagger v2', repo_id='SmilingWolf/wd-v1-4-moat-tagger-v2' ), + 'wd-v1-4-vit-tagger.v3': WaifuDiffusionInterrogator( + 'WD14 ViT v3', + repo_id='SmilingWolf/wd-vit-tagger-v3' + ), + 'wd-v1-4-convnext-tagger.v3': WaifuDiffusionInterrogator( + 'WD14 ConvNext v3', + repo_id='SmilingWolf/wd-convnext-tagger-v3' + ), + 'wd-v1-4-swinv2-tagger.v3': WaifuDiffusionInterrogator( + 'WD14 SwinV2 v3', + repo_id='SmilingWolf/wd-swinv2-tagger-v3' + ), 'mld-caformer.dec-5-97527': MLDanbooruInterrogator( 'ML-Danbooru Caformer dec-5-97527', repo_id='deepghs/ml-danbooru-onnx',