import gradio as gr import subprocess from .common_gui import ( get_folder_path, add_pre_postfix, find_replace, scriptdir, list_dirs, setup_environment, ) import os import sys from .custom_logging import setup_logging # Set up logging log = setup_logging() PYTHON = sys.executable def caption_images( caption_text: str, images_dir: str, overwrite: bool, caption_ext: str, prefix: str, postfix: str, find_text: str, replace_text: str, ): """ Captions images in a given directory with a given caption text. Args: caption_text (str): The text to be used as the caption. images_dir (str): The directory containing the images to be captioned. overwrite (bool): Whether to overwrite existing captions. caption_ext (str): The file extension for the caption files. prefix (str): Text to be added before the caption text. postfix (str): Text to be added after the caption text. find_text (str): Text to be replaced in the caption files. replace_text (str): Text to replace the found text in the caption files. Returns: None """ # Check if images_dir and caption_ext are provided missing_parameters = [] if not images_dir: missing_parameters.append("image directory") if not caption_ext: missing_parameters.append("caption file extension") if missing_parameters: log.info( "The following parameter(s) are missing: {}. " "Please provide these to proceed with captioning the images.".format(", ".join(missing_parameters)) ) return # Log the captioning process if caption_text: log.info(f"Captioning files in {images_dir} with {caption_text}...") # Build the command to run caption.py run_cmd = [ rf"{PYTHON}", rf"{scriptdir}/tools/caption.py", "--caption_text", caption_text, ] # Add optional flags to the command if overwrite: run_cmd.append("--overwrite") if caption_ext: run_cmd.append("--caption_file_ext") run_cmd.append(caption_ext) run_cmd.append(rf"{images_dir}") # Reconstruct the safe command string for display command_to_run = " ".join(run_cmd) log.info(f"Executing command: {command_to_run}") # Set the environment variable for the Python path env = setup_environment() # Run the command in the sd-scripts folder context subprocess.run(run_cmd, env=env, shell=False) # Check if overwrite option is enabled if overwrite: # Add prefix and postfix to caption files or find and replace text in caption files if prefix or postfix or find_text: # Add prefix and/or postfix to caption files add_pre_postfix( folder=images_dir, caption_file_ext=caption_ext, prefix=prefix, postfix=postfix, ) # Replace specified text in caption files if find and replace text is provided if find_text and replace_text: find_replace( folder_path=images_dir, caption_file_ext=caption_ext, search_text=find_text, replace_text=replace_text, ) else: # Show a message if modification is not possible without overwrite option enabled if prefix or postfix: log.info( 'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.' ) # Log the end of the captioning process log.info("Captioning done.") # Gradio UI def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None): """ Creates a Gradio tab for basic image captioning. Args: headless (bool, optional): If True, the GUI will be headless (no visible elements). Defaults to False. default_images_dir (str, optional): The default directory to use for image selection. If not provided, it defaults to the 'data' directory in the script directory. Returns: None """ from .common_gui import create_refresh_button # Set default images directory if not provided default_images_dir = ( default_images_dir if default_images_dir is not None else os.path.join(scriptdir, "data") ) current_images_dir = default_images_dir # Function to list directories def list_images_dirs(path): """ Lists directories within a specified path and updates the current image directory. Parameters: path (str): The directory path to list image directories from. Returns: list: A list of directories within the specified path. """ # Allows list_images_dirs to modify current_images_dir outside of this function nonlocal current_images_dir current_images_dir = path return list(list_dirs(path)) # Gradio tab for basic captioning with gr.Tab("Basic Captioning"): # Markdown description gr.Markdown( "This utility allows you to create simple caption files for each image in a folder." ) # Group and row for image folder selection with gr.Group(), gr.Row(): # Dropdown for image folder images_dir = gr.Dropdown( label="Image folder to caption (containing the images to caption)", choices=[""] + list_images_dirs(default_images_dir), value="", interactive=True, allow_custom_value=True, ) # Refresh button for image folder create_refresh_button( images_dir, lambda: None, lambda: {"choices": list_images_dirs(current_images_dir)}, "open_folder_small", ) # Button to open folder folder_button = gr.Button( "📂", elem_id="open_folder_small", elem_classes=["tool"], visible=(not headless), ) # Event handler for button click folder_button.click( get_folder_path, outputs=images_dir, show_progress=False, ) # Textbox for caption file extension caption_ext = gr.Dropdown( label="Caption file extension", choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, allow_custom_value=True, ) # Checkbox to overwrite existing captions overwrite = gr.Checkbox( label="Overwrite existing captions in folder", interactive=True, value=False, ) # Row for caption prefix and text with gr.Row(): # Textbox for caption prefix prefix = gr.Textbox( label="Prefix to add to caption", placeholder="(Optional)", interactive=True, ) # Textbox for caption text caption_text = gr.Textbox( label="Caption text", placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix or postfix.', interactive=True, lines=2, ) # Textbox for caption postfix postfix = gr.Textbox( label="Postfix to add to caption", placeholder="(Optional)", interactive=True, ) # Group and row for find and replace text with gr.Group(), gr.Row(): # Textbox for find text find_text = gr.Textbox( label="Find text", placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix or postfix.', interactive=True, lines=2, ) # Textbox for replace text replace_text = gr.Textbox( label="Replacement text", placeholder='e.g., "by some artist". Leave empty if you want to replace with nothing.', interactive=True, lines=2, ) # Button to caption images caption_button = gr.Button("Caption images") # Event handler for button click caption_button.click( caption_images, inputs=[ caption_text, images_dir, overwrite, caption_ext, prefix, postfix, find_text, replace_text, ], show_progress=False, ) # Event handler for dynamic update of dropdown choices images_dir.change( fn=lambda path: gr.Dropdown(choices=[""] + list_images_dirs(path)), inputs=images_dir, outputs=images_dir, show_progress=False, )