Merge pull request #312 from sALTaccount/master

Add replace underscores to WD14 tagger

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()
 

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,
         )