Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add append_captions feature to wd14 tagger #858

Merged
merged 2 commits into from
Oct 9, 2023
Merged

Add append_captions feature to wd14 tagger #858

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a4857fa
Add append_captions feature to wd14 tagger
a-l-e-x-d-s-9 Oct 5, 2023
9378da3
Fix comment
a-l-e-x-d-s-9 Oct 5, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 28 additions & 3 deletions finetune/tag_images_by_wd14_tagger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,12 +165,35 @@ def run_batch(path_imgs):
if len(character_tag_text) > 0:
character_tag_text = character_tag_text[2:]

caption_file = os.path.splitext(image_path)[0] + args.caption_extension

tag_text = ", ".join(combined_tags)

with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding="utf-8") as f:
if args.append_captions:
# Check if file exists
if os.path.exists(caption_file):

with open(caption_file, "rt", encoding="utf-8") as f:

# Read file and remove new lines
existing_content = f.read().strip("\n") # Remove newlines

# Split the content into tags and store them in a list
existing_tags = [tag.strip() for tag in existing_content.split(",") if tag.strip()]

# Check and remove repeating tags in tag_text
tag_text = ", ".join([tag for tag in combined_tags if tag not in existing_tags])

# If the file has content, prepend a comma to tag_text
if existing_content.strip() and tag_text:
tag_text = ", ".join(existing_tags) + ", " + tag_text


with open(caption_file, "wt", encoding="utf-8") as f:
f.write(tag_text + "\n")
if args.debug:
print(f"\n{image_path}:\n Character tags: {character_tag_text}\n General tags: {general_tag_text}")
print(
f"\n{image_path}:\n Character tags: {character_tag_text}\n General tags: {general_tag_text}")

# 読み込みの高速化のためにDataLoaderを使うオプション
if args.max_data_loader_n_workers is not None:
Expand Down Expand Up @@ -282,7 +305,9 @@ def setup_parser() -> argparse.ArgumentParser:
default="",
help="comma-separated list of undesired tags to remove from the output / 出力から除外したいタグのカンマ区切りのリスト",
)
parser.add_argument("--frequency_tags", action="store_true", help="Show frequency of tags for images / 画像ごとのタグの出現頻度を表示する")
parser.add_argument("--frequency_tags", action="store_true",
help="Show frequency of tags for images / 画像ごとのタグの出現頻度を表示する")
parser.add_argument("--append_captions", action="store_true", help="Append captions instead of overwriting")

return parser

Expand Down