StableDiffusion/KohyaSS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add new prefix and postfic for captions

Browse Source
This commit is contained in:
bmaltais 2022-12-19 11:43:29 -05:00
parent 6987f51b0a
commit 1d412726b3
3 changed files with 91 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
library/basic_caption_gui.py
View File

@ -1,37 +1,40 @@
import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import get_folder_path
from .common_gui import get_folder_path, add_pre_postfix
def caption_images(
caption_text_input, images_dir_input, overwrite_input, caption_file_ext
caption_text_input, images_dir_input, overwrite_input, caption_file_ext, prefix, postfix
):
# Check for caption_text_input
if caption_text_input == '':
msgbox('Caption text is missing...')
return
# Check for images_dir_input
if images_dir_input == '':
msgbox('Image folder is missing...')
return
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 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:
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)
# 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')
@ -46,22 +49,6 @@ def gradio_basic_caption_gui_tab():
gr.Markdown(
'This utility will allow the creation of simple caption files for each images in a folder.'
)
with gr.Row():
caption_text_input = gr.Textbox(
label='Caption text',
placeholder='Eg: , by some artist',
interactive=True,
)
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,
)
with gr.Row():
images_dir_input = gr.Textbox(
label='Image folder to caption',
@ -74,6 +61,33 @@ def gradio_basic_caption_gui_tab():
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(
@ -83,5 +97,6 @@ def gradio_basic_caption_gui_tab():
images_dir_input,
overwrite_input,
caption_file_ext,
prefix, postfix
],
)

25
library/blip_caption_gui.py
View File

@ -1,8 +1,8 @@
import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import get_folder_path
import os
from .common_gui import get_folder_path, add_pre_postfix
def caption_images(
train_data_dir,
@ -13,6 +13,8 @@ def caption_images(
max_length,
min_length,
beam_search,
prefix,
postfix
):
# Check for caption_text_input
# if caption_text_input == "":
@ -42,6 +44,9 @@ def caption_images(
# Run the command
subprocess.run(run_cmd)
# Add prefix and postfix
add_pre_postfix(folder=train_data_dir, caption_file_ext=caption_file_ext, prefix=prefix, postfix=postfix)
print('...captioning done')
@ -68,12 +73,24 @@ def gradio_blip_caption_gui_tab():
button_train_data_dir_input.click(
get_folder_path, outputs=train_data_dir
)
with gr.Row():
caption_file_ext = gr.Textbox(
label='Caption file extension',
placeholder='(Optional) Default: .caption',
interactive=True,
)
prefix = gr.Textbox(
label='Prefix to add to BLIP caption',
placeholder='(Optional)',
interactive=True,
)
postfix = gr.Textbox(
label='Postfix to add to BLIP caption',
placeholder='(Optional)',
interactive=True,
)
batch_size = gr.Number(
value=1, label='Batch size', interactive=True
@ -107,5 +124,7 @@ def gradio_blip_caption_gui_tab():
max_length,
min_length,
beam_search,
prefix,
postfix
],
)

18
library/common_gui.py
View File

@ -1,4 +1,5 @@
from tkinter import filedialog, Tk
import os
def get_file_path(file_path='', defaultextension='.json'):
current_file_path = file_path
@ -56,4 +57,19 @@ def get_saveasfile_path(file_path='', defaultextension='.json'):
print(file_path)
return file_path
return file_path
def add_pre_postfix(folder='', prefix='', postfix='', caption_file_ext='.caption'):
files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)]
if not prefix == '':
prefix = f'{prefix} '
if not postfix == '':
postfix = f' {postfix}'
for file in files:
with open(os.path.join(folder, file), 'r+') as f:
content = f.read()
content = content.rstrip()
f.seek(0,0)
f.write(f'{prefix}{content}{postfix}')
f.close()