Spaces:
Build error
Build error
File size: 3,365 Bytes
e69a9f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import get_folder_path
import os
def replace_underscore_with_space(folder_path, file_extension):
for file_name in os.listdir(folder_path):
if file_name.endswith(file_extension):
file_path = os.path.join(folder_path, file_name)
with open(file_path, 'r') as file:
file_content = file.read()
new_file_content = file_content.replace('_', ' ')
with open(file_path, 'w') as file:
file.write(new_file_content)
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...")
# return
# Check for images_dir_input
if train_data_dir == '':
msgbox('Image folder is missing...')
return
if caption_extension == '':
msgbox('Please provide an extension for the caption files.')
return
print(f'Captioning files in {train_data_dir}...')
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' --caption_extension="{caption_extension}"'
run_cmd += f' "{train_data_dir}"'
print(run_cmd)
# Run the command
if os.name == 'posix':
os.system(run_cmd)
else:
subprocess.run(run_cmd)
if replace_underscores:
replace_underscore_with_space(train_data_dir, caption_extension)
print('...captioning done')
###
# Gradio UI
###
def gradio_wd14_caption_gui_tab():
with gr.Tab('WD14 Captioning'):
gr.Markdown(
'This utility will use WD14 to caption files for each images in a folder.'
)
with gr.Row():
train_data_dir = gr.Textbox(
label='Image folder to caption',
placeholder='Directory containing the images to caption',
interactive=True,
)
button_train_data_dir_input = gr.Button(
'📂', elem_id='open_folder_small'
)
button_train_data_dir_input.click(
get_folder_path,
outputs=train_data_dir,
show_progress=False,
)
caption_extension = gr.Textbox(
label='Caption file extension',
placeholder='Extention for caption file. eg: .caption, .txt',
value='.txt',
interactive=True,
)
thresh = gr.Number(value=0.35, label='Threshold')
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,
replace_underscores,
],
show_progress=False,
)
|