diff --git a/.gitignore b/.gitignore index fd6106c..75124ed 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ __pycache__/ .vscode/ .venv/ +venv/ .env presets/ \ No newline at end of file diff --git a/README.md b/README.md index a57f8f1..9bd3cca 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,33 @@ ``` venv/bin/python -m venv venv -venv/bin/pip install -r requirements.txt +source venv/bin/activate +pip install -r requirements.txt + ``` ## usage ``` -venv/bin/python run.py image.jpg +usage: run.py [-h] [--dir DIR | --file FILE] [--threthold THRETHOLD] [--ext EXT] + +options: + -h, --help show this help message and exit + --dir DIR ディレクトリ内の画像すべてに予測を行う + --file FILE ファイルに対して予測を行う + --threthold THRETHOLD + 予測値の足切り確率(デフォルトは0.35) + --ext EXT dirの場合にキャプションファイルにつける拡張子 +``` + +single file + +``` +python3 run.py --file image.jpg +``` + +batch execution + +``` +python3 run.py --dir dir/dir ``` \ No newline at end of file diff --git a/run.py b/run.py index d8d5bed..f8d1564 100644 --- a/run.py +++ b/run.py @@ -4,21 +4,54 @@ import sys from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator from PIL import Image from pathlib import Path +import argparse +parser = argparse.ArgumentParser() + +group = parser.add_mutually_exclusive_group() +group.add_argument('--dir', help='ディレクトリ内の画像すべてに予測を行う') +group.add_argument('--file', help='ファイルに対して予測を行う') + +parser.add_argument( + '--threthold', + type=float, + default=0.35, + help='予測値の足切り確率(デフォルトは0.35)') +parser.add_argument( + '--ext', + default='.txt', + help='dirの場合にキャプションファイルにつける拡張子') + + +args = parser.parse_args() + +# 使用するモデル interrogator = WaifuDiffusionInterrogator( 'wd14-convnextv2-v2', repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2', revision='v2.0' ) -image_path = Path(sys.argv[1]) +def image_interrogate(image_path: Path): + """ + 画像パスから予測を行う + """ + im = Image.open(image_path) + result = interrogator.interrogate(im) + return Interrogator.postprocess_tags(result[1], threshold=0.35) -im = Image.open(image_path) -result = interrogator.interrogate(im) -tags = Interrogator.postprocess_tags(result[1], threshold=0.35) -print() +if args.dir: + d = Path(args.dir) + for f in d.iterdir(): + tags = image_interrogate(Path(f)) + + tags_str = ", ".join(tags.keys()) + with open(f.parent / f"{f.stem}{args.ext}", "w") as fp: + fp.write(tags_str) + +if args.file: + tags = image_interrogate(Path(args.file)) + tags_str = ", ".join(tags.keys()) + print(tags_str) -for k, v in tags.items(): - print(k, v) -print(result[0])