diff --git a/library/basic_caption_gui.py b/library/basic_caption_gui.py index 377a4b3..d4d187b 100644 --- a/library/basic_caption_gui.py +++ b/library/basic_caption_gui.py @@ -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 ], ) diff --git a/library/blip_caption_gui.py b/library/blip_caption_gui.py index fe33ca7..5678978 100644 --- a/library/blip_caption_gui.py +++ b/library/blip_caption_gui.py @@ -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 ], ) diff --git a/library/common_gui.py b/library/common_gui.py index 3979f77..8d68c58 100644 --- a/library/common_gui.py +++ b/library/common_gui.py @@ -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 \ No newline at end of file + 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() \ No newline at end of file