25 lines
513 B
Python
25 lines
513 B
Python
|
|
|
|
import sys
|
|
from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
|
|
interrogator = WaifuDiffusionInterrogator(
|
|
'wd14-convnextv2-v2',
|
|
repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2',
|
|
revision='v2.0'
|
|
)
|
|
|
|
image_path = Path(sys.argv[1])
|
|
|
|
im = Image.open(image_path)
|
|
result = interrogator.interrogate(im)
|
|
tags = Interrogator.postprocess_tags(result[1], threshold=0.35)
|
|
print()
|
|
|
|
for k, v in tags.items():
|
|
print(k, v)
|
|
print(result[0])
|
|
|