kohya_ss / library /group_images_gui.py
Ateras's picture
Upload folder using huggingface_hub
fe6327d
raw
history blame contribute delete
No virus
4.13 kB
import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import get_folder_path
import os
from library.custom_logging import setup_logging
# Set up logging
log = setup_logging()
PYTHON = 'python3' if os.name == 'posix' else './venv/Scripts/python.exe'
def group_images(
input_folder,
output_folder,
group_size,
include_subfolders,
do_not_copy_other_files,
generate_captions,
caption_ext
):
if input_folder == '':
msgbox('Input folder is missing...')
return
if output_folder == '':
msgbox('Please provide an output folder.')
return
log.info(f'Grouping images in {input_folder}...')
run_cmd = f'{PYTHON} "{os.path.join("tools","group_images.py")}"'
run_cmd += f' "{input_folder}"'
run_cmd += f' "{output_folder}"'
run_cmd += f' {(group_size)}'
if include_subfolders:
run_cmd += f' --include_subfolders'
if do_not_copy_other_files:
run_cmd += f' --do_not_copy_other_files'
if generate_captions:
run_cmd += f' --caption'
if caption_ext:
run_cmd += f' --caption_ext={caption_ext}'
log.info(run_cmd)
if os.name == 'posix':
os.system(run_cmd)
else:
subprocess.run(run_cmd)
log.info('...grouping done')
def gradio_group_images_gui_tab(headless=False):
with gr.Tab('Group Images'):
gr.Markdown('This utility will group images in a folder based on their aspect ratio.')
with gr.Row():
input_folder = gr.Textbox(
label='Input folder',
placeholder='Directory containing the images to group',
interactive=True,
)
button_input_folder = gr.Button(
'πŸ“‚', elem_id='open_folder_small', visible=(not headless)
)
button_input_folder.click(
get_folder_path,
outputs=input_folder,
show_progress=False,
)
output_folder = gr.Textbox(
label='Output folder',
placeholder='Directory where the grouped images will be stored',
interactive=True,
)
button_output_folder = gr.Button(
'πŸ“‚', elem_id='open_folder_small', visible=(not headless)
)
button_output_folder.click(
get_folder_path,
outputs=output_folder,
show_progress=False,
)
with gr.Row():
group_size = gr.Slider(
label='Group size',
info='Number of images to group together',
value='4',
minimum=1, maximum=64, step=1,
interactive=True,
)
include_subfolders = gr.Checkbox(
label='Include Subfolders',
value=False,
info='Include images in subfolders as well',
)
do_not_copy_other_files = gr.Checkbox(
label='Do not copy other files',
value=False,
info='Do not copy other files in the input folder to the output folder',
)
generate_captions = gr.Checkbox(
label='Generate Captions',
value=False,
info='Generate caption files for the grouped images based on their folder name',
)
caption_ext = gr.Textbox(
label='Caption Extension',
placeholder='Caption file extension (e.g., .txt)',
value='.txt',
interactive=True,
)
group_images_button = gr.Button('Group images')
group_images_button.click(
group_images,
inputs=[
input_folder,
output_folder,
group_size,
include_subfolders,
do_not_copy_other_files,
generate_captions,
caption_ext,
],
show_progress=False,
)