diff --git a/README.md b/README.md index b79b3dc..190c594 100644 --- a/README.md +++ b/README.md @@ -13,18 +13,21 @@ pip install -r requirements.txt ## usage ``` -usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRETHOLD] [--ext EXT] +usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--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}] + [--overwrite OVERWRITE] options: -h, --help show this help message and exit - --dir DIR ディレクトリ内の画像すべてに予測を行う - --file FILE ファイルに対して予測を行う - --threshold THRETHOLD - 予測値の足切り確率(デフォルトは0.35) - --ext EXT dirの場合にキャプションファイルにつける拡張子 + --dir DIR Predictions for all images in the directory + --file FILE Predictions for one file + --threshold THRESHOLD + Prediction threshold (default is 0.35) + --ext EXT Extension to add to caption file in case of dir option (default is .txt) --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} - 予測に使用するモデル名 + modelname to use for prediction (default is wd14-convnextv2.v1) + --overwrite OVERWRITE + Overwrite the file if it exists (default is True) ``` single file diff --git a/run.py b/run.py index d491879..b080dc0 100644 --- a/run.py +++ b/run.py @@ -11,32 +11,36 @@ from tagger.interrogators import interrogators parser = argparse.ArgumentParser() group = parser.add_mutually_exclusive_group(required=True) -group.add_argument('--dir', help='ディレクトリ内の画像すべてに予測を行う') -group.add_argument('--file', help='ファイルに対して予測を行う') +group.add_argument('--dir', help='Predictions for all images in the directory') +group.add_argument('--file', help='Predictions for one file') parser.add_argument( '--threshold', type=float, default=0.35, - help='予測値の足切り確率(デフォルトは0.35)') + help='Prediction threshold (default is 0.35)') parser.add_argument( '--ext', default='.txt', - help='dirの場合にキャプションファイルにつける拡張子') + help='Extension to add to caption file in case of dir option (default is .txt)') parser.add_argument( '--model', default='wd14-convnextv2.v1', choices=list(interrogators.keys()), - help='予測に使用するモデル名') - + help='modelname to use for prediction (default is wd14-convnextv2.v1)') +parser.add_argument( + '--overwrite', + default=True, + type=bool, + help='Overwrite the file if it exists (default is True)') args = parser.parse_args() -# 使用するinterrogatorを読み出す +# get interrogator configs interrogator = interrogators[args.model] def image_interrogate(image_path: Path): """ - 画像パスから予測を行う + Predictions from a image path """ im = Image.open(image_path) result = interrogator.interrogate(im) @@ -45,7 +49,7 @@ def image_interrogate(image_path: Path): if args.dir: d = Path(args.dir) for f in d.iterdir(): - if not f.is_file() or f.suffix not in ['.png', '.jpg', '.webp']: + if not f.is_file() or f.suffix not in ['.png', '.jpg', '.jpeg', '.webp']: continue image_path = Path(f) print('processing:', image_path)