File size: 4,470 Bytes
e69a9f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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,
            )