Merge pull request #312 from sALTaccount/master
Add replace underscores to WD14 tagger
This commit is contained in:
commit
a57cdd5d42
@ -125,7 +125,7 @@ def main(args):
|
|||||||
tag_text = ""
|
tag_text = ""
|
||||||
for i, p in enumerate(prob[4:]): # numpyとか使うのが良いけど、まあそれほど数も多くないのでループで
|
for i, p in enumerate(prob[4:]): # numpyとか使うのが良いけど、まあそれほど数も多くないのでループで
|
||||||
if p >= args.thresh and i < len(tags):
|
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:
|
if len(tag_text) > 0:
|
||||||
tag_text = tag_text[2:] # 最初の ", " を消す
|
tag_text = tag_text[2:] # 最初の ", " を消す
|
||||||
@ -190,6 +190,7 @@ if __name__ == '__main__':
|
|||||||
help="extension of caption file (for backward compatibility) / 出力されるキャプションファイルの拡張子(スペルミスしていたのを残してあります)")
|
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("--caption_extension", type=str, default=".txt", help="extension of caption file / 出力されるキャプションファイルの拡張子")
|
||||||
parser.add_argument("--debug", action="store_true", help="debug mode")
|
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()
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
@ -5,7 +5,7 @@ from .common_gui import get_folder_path
|
|||||||
import os
|
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
|
# Check for caption_text_input
|
||||||
# if caption_text_input == "":
|
# if caption_text_input == "":
|
||||||
# msgbox("Caption text is missing...")
|
# 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'accelerate launch "./finetune/tag_images_by_wd14_tagger.py"'
|
||||||
run_cmd += f' --batch_size="{int(batch_size)}"'
|
run_cmd += f' --batch_size="{int(batch_size)}"'
|
||||||
run_cmd += f' --thresh="{thresh}"'
|
run_cmd += f' --thresh="{thresh}"'
|
||||||
|
run_cmd += f' --replace_underscores' if replace_underscores else ''
|
||||||
if caption_extension != '':
|
if caption_extension != '':
|
||||||
run_cmd += f' --caption_extension="{caption_extension}"'
|
run_cmd += f' --caption_extension="{caption_extension}"'
|
||||||
run_cmd += f' "{train_data_dir}"'
|
run_cmd += f' "{train_data_dir}"'
|
||||||
@ -75,11 +76,17 @@ def gradio_wd14_caption_gui_tab():
|
|||||||
batch_size = gr.Number(
|
batch_size = gr.Number(
|
||||||
value=1, label='Batch size', interactive=True
|
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 = gr.Button('Caption images')
|
||||||
|
|
||||||
caption_button.click(
|
caption_button.click(
|
||||||
caption_images,
|
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,
|
show_progress=False,
|
||||||
)
|
)
|
||||||
|
Loading…
Reference in New Issue
Block a user