KohyaSS/library/basic_caption_gui.py

import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import get_folder_path, add_pre_postfix, find_replace
import os


def caption_images(
    caption_text,
    images_dir,
    overwrite,
    caption_ext,
    prefix,
    postfix,
    find_text,
    replace_text,
):
    # Check for images_dir
    if not images_dir:
        msgbox('Image folder is missing...')
        return

    if not caption_ext:
        msgbox('Please provide an extension for the caption files.')
        return

    if caption_text:
        print(f'Captioning files in {images_dir} with {caption_text}...')
        run_cmd = f'python "tools/caption.py"'
        run_cmd += f' --caption_text="{caption_text}"'
        if overwrite:
            run_cmd += f' --overwrite'
        if caption_ext:
            run_cmd += f' --caption_file_ext="{caption_ext}"'
        run_cmd += f' "{images_dir}"'

        print(run_cmd)

        # Run the command
        if os.name == 'posix':
            os.system(run_cmd)
        else:
            subprocess.run(run_cmd)

    if overwrite:
        if prefix or postfix:
            # Add prefix and postfix
            add_pre_postfix(
                folder=images_dir,
                caption_file_ext=caption_ext,
                prefix=prefix,
                postfix=postfix,
            )
        if find_text:
            find_replace(
                folder_path=images_dir,
                caption_file_ext=caption_ext,
                search_text=find_text,
                replace_text=replace_text,
            )
    else:
        if prefix or postfix:
            msgbox(
                'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected...'
            )

    print('...captioning done')


# Gradio UI
def gradio_basic_caption_gui_tab():
    with gr.Tab('Basic Captioning'):
        gr.Markdown(
            'This utility will allow the creation of simple caption files for each image in a folder.'
        )
        with gr.Row():
            images_dir = gr.Textbox(
                label='Image folder to caption',
                placeholder='Directory containing the images to caption',
                interactive=True,
            )
            folder_button = gr.Button('📂', elem_id='open_folder_small')
            folder_button.click(
                get_folder_path,
                outputs=images_dir,
                show_progress=False,
            )
            caption_ext = gr.Textbox(
                label='Caption file extension',
                placeholder='Extension for caption file. eg: .caption, .txt',
                value='.txt',
                interactive=True,
            )
            overwrite = gr.Checkbox(
                label='Overwrite existing captions in folder',
                interactive=True,
                value=False,
            )
        with gr.Row():
            prefix = gr.Textbox(
                label='Prefix to add to caption',
                placeholder='(Optional)',
                interactive=True,
            )
            caption_text = gr.Textbox(
                label='Caption text',
                placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix',
                interactive=True,
            )
            postfix = gr.Textbox(
                label='Postfix to add to caption',
                placeholder='(Optional)',
                interactive=True,
            )
        with gr.Row():
            find_text = gr.Textbox(
                label='Find text',
                placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix',
                interactive=True,
            )
            replace_text = 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')
            caption_button.click(
                caption_images,
                inputs=[
                    caption_text,
                    images_dir,
                    overwrite,
                    caption_ext,
                    prefix,
                    postfix,
                    find_text,
                    replace_text,
                ],
                show_progress=False,
            )