diff --git a/finetune/tag_images_by_wd14_tagger.py b/finetune/tag_images_by_wd14_tagger.py index 609b8c5..933117c 100644 --- a/finetune/tag_images_by_wd14_tagger.py +++ b/finetune/tag_images_by_wd14_tagger.py @@ -125,7 +125,7 @@ def main(args): tag_text = "" for i, p in enumerate(prob[4:]): # numpyとか使うのが良いけど、まあそれほど数も多くないのでループで if p >= args.thresh and i < len(tags): - tag_text += ", " + tags[i] + tag_text += ", " + (tags[i].replace("_", " ") if args.replace_underscores else tags[i]) if len(tag_text) > 0: tag_text = tag_text[2:] # 最初の ", " を消す @@ -190,6 +190,7 @@ if __name__ == '__main__': help="extension of caption file (for backward compatibility) / 出力されるキャプションファイルの拡張子(スペルミスしていたのを残してあります)") parser.add_argument("--caption_extension", type=str, default=".txt", help="extension of caption file / 出力されるキャプションファイルの拡張子") parser.add_argument("--debug", action="store_true", help="debug mode") + parser.add_argument("--replace_underscores", action="store_true", help="replace underscores in tags with spaces / タグのアンダースコアをスペースに置き換える") args = parser.parse_args() diff --git a/library/wd14_caption_gui.py b/library/wd14_caption_gui.py index 89de66d..99d22d8 100644 --- a/library/wd14_caption_gui.py +++ b/library/wd14_caption_gui.py @@ -5,7 +5,7 @@ from .common_gui import get_folder_path import os -def caption_images(train_data_dir, caption_extension, batch_size, thresh): +def caption_images(train_data_dir, caption_extension, batch_size, thresh, replace_underscores): # Check for caption_text_input # if caption_text_input == "": # msgbox("Caption text is missing...") @@ -24,6 +24,7 @@ def caption_images(train_data_dir, caption_extension, batch_size, thresh): run_cmd = f'accelerate launch "./finetune/tag_images_by_wd14_tagger.py"' run_cmd += f' --batch_size="{int(batch_size)}"' run_cmd += f' --thresh="{thresh}"' + run_cmd += f' --replace_underscores' if replace_underscores else '' if caption_extension != '': run_cmd += f' --caption_extension="{caption_extension}"' run_cmd += f' "{train_data_dir}"' @@ -75,11 +76,17 @@ def gradio_wd14_caption_gui_tab(): batch_size = gr.Number( value=1, label='Batch size', interactive=True ) + + replace_underscores = gr.Checkbox( + label='Replace underscores in filenames with spaces', + value=False, + interactive=True, + ) caption_button = gr.Button('Caption images') caption_button.click( caption_images, - inputs=[train_data_dir, caption_extension, batch_size, thresh], + inputs=[train_data_dir, caption_extension, batch_size, thresh, replace_underscores], show_progress=False, )