diff --git a/README.md b/README.md index 0a231ff..33ae9a6 100644 --- a/README.md +++ b/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) ``` diff --git a/run.py b/run.py index 07e86dc..92aa2ce 100644 --- a/run.py +++ b/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) diff --git a/tagger/interrogator.py b/tagger/interrogator.py index c48caf0..4f1d42a 100644 --- a/tagger/interrogator.py +++ b/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,