Spaces:
Sleeping
Sleeping
File size: 2,865 Bytes
b78b1f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import argparse
import re
from pathlib import Path
def load_json_dict(path: str):
import json
from pathlib import Path
dict = {}
if not Path(path).exists(): return dict
try:
with open(path, encoding='utf-8') as f:
dict = json.load(f)
except Exception:
print(f"Failed to open dictionary file: {path}")
return dict
return dict
danbooru_ja_dict = load_json_dict('danbooru_ja_dict.json')
def danbooru_tags_to_jas(tags: list[str]):
from rapidfuzz.process import extractOne
from rapidfuzz.utils import default_process
A1111_SPECIAL_SYNTAX_RE = re.compile(r"\s*<[^>]+>\s*")
keys = list(danbooru_ja_dict.keys())
jas = []
for tag in tags:
tag = str(tag).strip()
if A1111_SPECIAL_SYNTAX_RE.fullmatch(tag): continue
s = default_process(str(tag))
e1 = extractOne(s, keys, processor=default_process, score_cutoff=90.0)
if e1:
jas.extend(danbooru_ja_dict[e1[0]].copy())
return jas
def danbooru_to_ja(input_tag, input_file, output_file, is_append):
if input_file and Path(input_file).exists():
try:
with open(input_file, 'r', encoding='utf-8') as f:
input_tag = f.read()
except Exception:
print(f"Failed to open input file: {input_file}")
tags = [tag.strip() for tag in input_tag.split(",")] if input_tag else []
ja_tags = danbooru_tags_to_jas(tags)
output_tags = tags + ja_tags if is_append else ja_tags
output_tag = ", ".join(output_tags)
if output_file:
try:
with open(output_file, mode='w', encoding="utf-8") as f:
f.write(output_tag)
except Exception:
print(f"Failed to write output file: {output_file}")
else:
print(output_tag)
return output_tag
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tags", default=None, type=str, required=False, help="Input tags.")
parser.add_argument("--file", default=None, type=str, required=False, help="Input tags from a text file.")
parser.add_argument("--out", default=None, type=str, help="Output to text file.")
parser.add_argument("--append", default=False, type=bool, help="Whether the output contains the input tags or not.")
args = parser.parse_args()
assert (args.tags, args.file) != (None, None), "Must provide --tags or --file!"
danbooru_to_ja(args.tags, args.file, args.out, args.append)
# Usage:
# python danbooru_to_ja.py --tags "1girl, oomuro sakurako, solo, sitting, starry sky"
# python danbooru_to_ja.py --file inputtag.txt
# python danbooru_to_ja.py --file inputtag.txt --append True
# Datasets: https://huggingface.co/datasets/p1atdev/danbooru-ja-tag-pair-20240715 |