1f1dd5c4de
shuffle_caption, save_state, resume, prior_loss_weight, Fix issue with config open and save
116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
import gradio as gr
|
|
from easygui import msgbox
|
|
import subprocess
|
|
from .common_gui import get_folder_path, add_pre_postfix
|
|
|
|
|
|
def caption_images(
|
|
caption_text_input,
|
|
images_dir_input,
|
|
overwrite_input,
|
|
caption_file_ext,
|
|
prefix,
|
|
postfix,
|
|
):
|
|
# Check for images_dir_input
|
|
if images_dir_input == '':
|
|
msgbox('Image folder is missing...')
|
|
return
|
|
|
|
if not caption_text_input == '':
|
|
print(
|
|
f'Captioning files in {images_dir_input} with {caption_text_input}...'
|
|
)
|
|
run_cmd = f'python "tools/caption.py"'
|
|
run_cmd += f' --caption_text="{caption_text_input}"'
|
|
if overwrite_input:
|
|
run_cmd += f' --overwrite'
|
|
if caption_file_ext != '':
|
|
run_cmd += f' --caption_file_ext="{caption_file_ext}"'
|
|
run_cmd += f' "{images_dir_input}"'
|
|
|
|
print(run_cmd)
|
|
|
|
# Run the command
|
|
subprocess.run(run_cmd)
|
|
|
|
if overwrite_input:
|
|
# Add prefix and postfix
|
|
add_pre_postfix(
|
|
folder=images_dir_input,
|
|
caption_file_ext=caption_file_ext,
|
|
prefix=prefix,
|
|
postfix=postfix,
|
|
)
|
|
else:
|
|
if not prefix == '' or not postfix == '':
|
|
msgbox(
|
|
'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected...'
|
|
)
|
|
|
|
print('...captioning done')
|
|
|
|
|
|
###
|
|
# Gradio UI
|
|
###
|
|
|
|
|
|
def gradio_basic_caption_gui_tab():
|
|
with gr.Tab('Basic Captioning'):
|
|
gr.Markdown(
|
|
'This utility will allow the creation of simple caption files for each images in a folder.'
|
|
)
|
|
with gr.Row():
|
|
images_dir_input = gr.Textbox(
|
|
label='Image folder to caption',
|
|
placeholder='Directory containing the images to caption',
|
|
interactive=True,
|
|
)
|
|
button_images_dir_input = gr.Button(
|
|
'📂', elem_id='open_folder_small'
|
|
)
|
|
button_images_dir_input.click(
|
|
get_folder_path, outputs=images_dir_input
|
|
)
|
|
with gr.Row():
|
|
prefix = gr.Textbox(
|
|
label='Prefix to add to txt caption',
|
|
placeholder='(Optional)',
|
|
interactive=True,
|
|
)
|
|
caption_text_input = gr.Textbox(
|
|
label='Caption text',
|
|
placeholder='Eg: , by some artist. Leave empti if you just want to add pre or postfix',
|
|
interactive=True,
|
|
)
|
|
postfix = gr.Textbox(
|
|
label='Postfix to add to txt caption',
|
|
placeholder='(Optional)',
|
|
interactive=True,
|
|
)
|
|
with gr.Row():
|
|
overwrite_input = gr.Checkbox(
|
|
label='Overwrite existing captions in folder',
|
|
interactive=True,
|
|
value=False,
|
|
)
|
|
caption_file_ext = gr.Textbox(
|
|
label='Caption file extension',
|
|
placeholder='(Optional) Default: .caption',
|
|
interactive=True,
|
|
)
|
|
caption_button = gr.Button('Caption images')
|
|
|
|
caption_button.click(
|
|
caption_images,
|
|
inputs=[
|
|
caption_text_input,
|
|
images_dir_input,
|
|
overwrite_input,
|
|
caption_file_ext,
|
|
prefix,
|
|
postfix,
|
|
],
|
|
)
|