diff --git a/library/basic_caption_gui.py b/library/basic_caption_gui.py index 3c1d719..550d757 100644 --- a/library/basic_caption_gui.py +++ b/library/basic_caption_gui.py @@ -1,7 +1,7 @@ import gradio as gr from easygui import msgbox import subprocess -from .common_gui import get_folder_path, add_pre_postfix +from .common_gui import get_folder_path, add_pre_postfix, find_replace def caption_images( @@ -10,7 +10,7 @@ def caption_images( overwrite_input, caption_file_ext, prefix, - postfix, + postfix, find, replace ): # Check for images_dir_input if images_dir_input == '': @@ -35,13 +35,21 @@ def caption_images( 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, - ) + if not prefix=='' or not postfix=='': + # Add prefix and postfix + add_pre_postfix( + folder=images_dir_input, + caption_file_ext=caption_file_ext, + prefix=prefix, + postfix=postfix, + ) + if not find=='': + find_replace( + folder=images_dir_input, + caption_file_ext=caption_file_ext, + find=find, + replace=replace, + ) else: if not prefix == '' or not postfix == '': msgbox( @@ -73,6 +81,16 @@ def gradio_basic_caption_gui_tab(): button_images_dir_input.click( get_folder_path, outputs=images_dir_input ) + caption_file_ext = gr.Textbox( + label='Caption file extension', + placeholder='(Optional) Default: .caption', + interactive=True, + ) + overwrite_input = gr.Checkbox( + label='Overwrite existing captions in folder', + interactive=True, + value=False, + ) with gr.Row(): prefix = gr.Textbox( label='Prefix to add to caption', @@ -81,7 +99,7 @@ def gradio_basic_caption_gui_tab(): ) caption_text_input = gr.Textbox( label='Caption text', - placeholder='Eg: , by some artist. Leave empti if you just want to add pre or postfix', + placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix', interactive=True, ) postfix = gr.Textbox( @@ -90,14 +108,14 @@ def gradio_basic_caption_gui_tab(): interactive=True, ) with gr.Row(): - overwrite_input = gr.Checkbox( - label='Overwrite existing captions in folder', + find = gr.Textbox( + label='Find text', + placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix', interactive=True, - value=False, ) - caption_file_ext = gr.Textbox( - label='Caption file extension', - placeholder='(Optional) Default: .caption', + replace = 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') @@ -111,5 +129,6 @@ def gradio_basic_caption_gui_tab(): caption_file_ext, prefix, postfix, + find, replace ], ) diff --git a/library/common_gui.py b/library/common_gui.py index f54267b..ff7b2c5 100644 --- a/library/common_gui.py +++ b/library/common_gui.py @@ -7,6 +7,15 @@ def get_dir_and_file(file_path): dir_path, file_name = os.path.split(file_path) return (dir_path, file_name) +def has_ext_files(directory, extension): + # Iterate through all the files in the directory + for file in os.listdir(directory): + # If the file name ends with extension, return True + if file.endswith(extension): + return True + # If no extension files were found, return False + return False + def get_file_path(file_path='', defaultextension='.json', extension_name='Config files'): current_file_path = file_path # print(f'current file path: {current_file_path}') @@ -127,13 +136,13 @@ def get_saveasfilename_path(file_path='', extensions='*', extension_name='Config def add_pre_postfix( folder='', prefix='', postfix='', caption_file_ext='.caption' ): + if not has_ext_files(folder, caption_file_ext): + msgbox(f'No files with extension {caption_file_ext} were found in {folder}...') + return + if prefix == '' and postfix == '': return - # set caption extention to default in case it was not provided - if caption_file_ext == '': - caption_file_ext = '.caption' - files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)] if not prefix == '': prefix = f'{prefix} ' @@ -147,6 +156,27 @@ def add_pre_postfix( f.seek(0, 0) f.write(f'{prefix}{content}{postfix}') f.close() + +def find_replace( + folder='', caption_file_ext='.caption', find='', replace='' +): + print('Running caption find/replace') + if not has_ext_files(folder, caption_file_ext): + msgbox(f'No files with extension {caption_file_ext} were found in {folder}...') + return + + if find == '': + return + + files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)] + for file in files: + with open(os.path.join(folder, file), 'r') as f: + content = f.read() + f.close + content = content.replace(find, replace) + with open(os.path.join(folder, file), 'w') as f: + f.write(content) + f.close() def color_aug_changed(color_aug): if color_aug: