7b5639cff5
This is a massive WIP and should not be trusted or used right now. However, major milestones have been crossed. Both message boxes and file dialogs are now properly subprocessed and work on macOS. I think by extension, it may work on runpod environments as well, but that remains to be tested.
142 lines
4.4 KiB
Python
142 lines
4.4 KiB
Python
import os
|
|
import subprocess
|
|
|
|
import gradio as gr
|
|
|
|
from .common_gui import get_folder_path, add_pre_postfix, find_replace
|
|
|
|
|
|
def caption_images(
|
|
caption_text,
|
|
images_dir,
|
|
overwrite,
|
|
caption_ext,
|
|
prefix,
|
|
postfix,
|
|
find_text,
|
|
replace_text,
|
|
):
|
|
# Check for images_dir
|
|
if not images_dir:
|
|
show_message_box('Image folder is missing...')
|
|
return
|
|
|
|
if not caption_ext:
|
|
show_message_box('Please provide an extension for the caption files.')
|
|
return
|
|
|
|
if caption_text:
|
|
print(f'Captioning files in {images_dir} with {caption_text}...')
|
|
run_cmd = f'python "tools/caption.py"'
|
|
run_cmd += f' --caption_text="{caption_text}"'
|
|
if overwrite:
|
|
run_cmd += f' --overwrite'
|
|
if caption_ext:
|
|
run_cmd += f' --caption_file_ext="{caption_ext}"'
|
|
run_cmd += f' "{images_dir}"'
|
|
|
|
print(run_cmd)
|
|
|
|
# Run the command
|
|
if os.name == 'posix':
|
|
os.system(run_cmd)
|
|
else:
|
|
subprocess.run(run_cmd)
|
|
|
|
if overwrite:
|
|
if prefix or postfix:
|
|
# Add prefix and postfix
|
|
add_pre_postfix(
|
|
folder=images_dir,
|
|
caption_file_ext=caption_ext,
|
|
prefix=prefix,
|
|
postfix=postfix,
|
|
)
|
|
if find_text:
|
|
find_replace(
|
|
folder_path=images_dir,
|
|
caption_file_ext=caption_ext,
|
|
search_text=find_text,
|
|
replace_text=replace_text,
|
|
)
|
|
else:
|
|
if prefix or postfix:
|
|
show_message_box(
|
|
'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 image in a folder.'
|
|
)
|
|
with gr.Row():
|
|
images_dir = gr.Textbox(
|
|
label='Image folder to caption',
|
|
placeholder='Directory containing the images to caption',
|
|
interactive=True,
|
|
)
|
|
folder_button = gr.Button('📂', elem_id='open_folder_small')
|
|
folder_button.click(
|
|
get_folder_path,
|
|
outputs=images_dir,
|
|
show_progress=False,
|
|
)
|
|
caption_ext = gr.Textbox(
|
|
label='Caption file extension',
|
|
placeholder='Extension for caption file. eg: .caption, .txt',
|
|
value='.txt',
|
|
interactive=True,
|
|
)
|
|
overwrite = gr.Checkbox(
|
|
label='Overwrite existing captions in folder',
|
|
interactive=True,
|
|
value=False,
|
|
)
|
|
with gr.Row():
|
|
prefix = gr.Textbox(
|
|
label='Prefix to add to caption',
|
|
placeholder='(Optional)',
|
|
interactive=True,
|
|
)
|
|
caption_text = gr.Textbox(
|
|
label='Caption text',
|
|
placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix',
|
|
interactive=True,
|
|
)
|
|
postfix = gr.Textbox(
|
|
label='Postfix to add to caption',
|
|
placeholder='(Optional)',
|
|
interactive=True,
|
|
)
|
|
with gr.Row():
|
|
find_text = gr.Textbox(
|
|
label='Find text',
|
|
placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix',
|
|
interactive=True,
|
|
)
|
|
replace_text = gr.Textbox(
|
|
label='Replacement text',
|
|
placeholder='Eg: , by some artist. Leave empty if you just want to replace with nothing',
|
|
interactive=True,
|
|
)
|
|
caption_button = gr.Button('Caption images')
|
|
caption_button.click(
|
|
caption_images,
|
|
inputs=[
|
|
caption_text,
|
|
images_dir,
|
|
overwrite,
|
|
caption_ext,
|
|
prefix,
|
|
postfix,
|
|
find_text,
|
|
replace_text,
|
|
],
|
|
show_progress=False,
|
|
)
|