Add exclude-tag option.

README.md

@@ -12,7 +12,7 @@ pip install -r requirements.txt
 ## usage
 
 ```
-usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] [--recursive]
+usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] [--recursive] [--exclude-tag t1,t2,t3]
               [--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:
@@ -26,6 +26,8 @@ options:
   --cpu                 Use CPU only
   --rawtag              Use the raw output of the model
   --recursive           Enable recursive file search
+  --exclude-tag t1,t2,t3
+                        Specify tags to exclude (Need comma-separated list)
   --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)
 ```

run.py

@@ -1,4 +1,4 @@
-from typing import Generator
+from typing import Generator, Iterable
 from tagger.interrogator import Interrogator
 from PIL import Image
 from pathlib import Path
@@ -37,6 +37,12 @@ parser.add_argument(
     '--recursive',
     action='store_true',
     help='Enable recursive file search')
+parser.add_argument(
+    '--exclude-tag',
+    dest='exclude_tags',
+    action='append',
+    metavar='t1,t2,t3',
+    help='Specify tags to exclude (Need comma-separated list)')
 parser.add_argument(
     '--model',
     default='wd14-convnextv2.v1',
@@ -50,16 +56,35 @@ interrogator = interrogators[args.model]
 if args.cpu:
     interrogator.use_cpu()
 
-def image_interrogate(image_path: Path):
+def parse_exclude_tags() -> set[str]:
+    if args.exclude_tags is None:
+        return set()
+
+    tags = []
+    for str in args.exclude_tags:
+        for tag in str.split(','):
+            tags.append(tag.strip())
+
+    # reverse escape (nai tag to danbooru tag)
+    reverse_escaped_tags = []
+    for tag in tags:
+        tag = tag.replace(' ', '_').replace('\(', '(').replace('\)', ')')
+        reverse_escaped_tags.append(tag)
+    return set([*tags, *reverse_escaped_tags])  # reduce duplicates
+
+def image_interrogate(image_path: Path, tag_escape: bool, exclude_tags: Iterable[str]) -> dict[str, float]:
     """
     Predictions from a image path
     """
     im = Image.open(image_path)
     result = interrogator.interrogate(im)
 
-    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)
+    return Interrogator.postprocess_tags(
+        result[1],
+        threshold=args.threshold,
+        escape_tag=tag_escape,
+        replace_underscore=tag_escape,
+        exclude_tags=exclude_tags)
 
 def explore_image_files(folder_path: Path) -> Generator[Path, None, None]:
     """
@@ -82,7 +107,7 @@ if args.dir:
             continue
 
         print('processing:', image_path)
-        tags = image_interrogate(image_path)
+        tags = image_interrogate(image_path, not args.rawtag, parse_exclude_tags())
 
         tags_str = ', '.join(tags.keys())
 
@@ -90,7 +115,7 @@ if args.dir:
             fp.write(tags_str)
 
 if args.file:
-    tags = image_interrogate(Path(args.file))
+    tags = image_interrogate(Path(args.file), not args.rawtag, parse_exclude_tags())
     print()
     tags_str = ', '.join(tags.keys())
     print(tags_str)

tagger/interrogator.py

@@ -2,8 +2,7 @@ import os
 import pandas as pd
 import numpy as np
 
-from typing import Tuple, List, Dict
-from io import BytesIO
+from typing import Iterable, Tuple, List, Dict
 from PIL import Image
 
 from pathlib import Path
@@ -24,7 +23,7 @@ class Interrogator:
         tags: Dict[str, float],
         threshold=0.35,
         additional_tags: List[str] = [],
-        exclude_tags: List[str] = [],
+        exclude_tags: Iterable[str] = [],
         sort_by_alphabetical_order=False,
         add_confident_as_weight=False,
         replace_underscore=False,