Add wd14 v3 models.

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
 ```

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)
 
 

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',