pi-ux-ce/wd14-tagger-standalone
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #9 from corkborg/support-recursive

Browse Source 
Support recursive.
This commit is contained in:
corkborg
2024-06-03 00:49:09 +09:00
committed by GitHub
parent 648abba071 4c421455f9
commit c502aff885
2 changed files with 23 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2 -2
View File
@@ -1,7 +1,6 @@
forked from [https://github.com/picobyte/stable-diffusion-webui-wd14-tagger](https://github.com/picobyte/stable-diffusion-webui-wd14-tagger)
## install
```
@@ -13,7 +12,7 @@ pip install -r requirements.txt
## usage
```
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag]
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] [--recursive]
[--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 +25,7 @@ options:
--overwrite Overwrite caption file if it exists
--cpu Use CPU only
--rawtag Use the raw output of the model
--recursive Enable recursive file search
--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
+21 -10
View File
@@ -1,4 +1,5 @@
from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
from typing import Generator
from tagger.interrogator import Interrogator
from PIL import Image
from pathlib import Path
import argparse
@@ -32,6 +33,10 @@ parser.add_argument(
'--rawtag',
action='store_true',
help='Use the raw output of the model')
parser.add_argument(
'--recursive',
action='store_true',
help='Enable recursive file search')
parser.add_argument(
'--model',
default='wd14-convnextv2.v1',
@@ -56,17 +61,23 @@ def image_interrogate(image_path: Path):
return Interrogator.postprocess_tags(result[1], threshold=args.threshold)
return Interrogator.postprocess_tags(result[1], threshold=args.threshold, escape_tag=True, replace_underscore=True)
def explore_image_files(folder_path: Path) -> Generator[Path, None, None]:
"""
Explore files by folder path
"""
for path in folder_path.iterdir():
if path.is_file() and path.suffix in ['.png', '.jpg', '.jpeg', '.webp']:
yield path
elif args.recursive and path.is_dir():
yield from explore_image_files(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', '.jpeg', '.webp']:
continue
root_path = Path(args.dir)
for image_path in explore_image_files(root_path):
caption_path = image_path.parent / f'{image_path.stem}{args.ext}'
image_path = Path(f)
caption_path = image_path.parent / f'{f.stem}{args.ext}'
if caption_path.is_file() and not args.overwrite:
# skip if file exists
if caption_path.is_file() and not args.overwrite:
# skip if caption exists
print('skip:', image_path)
continue