Add gui for group_images

kohya_gui.py

@@ -118,12 +118,13 @@ def UI(**kwargs):
                 enable_copy_info_button=True,
                 headless=headless,
             )
-            gradio_extract_dylora_tab(headless=headless)
-            gradio_extract_lora_tab(headless=headless)
-            gradio_extract_lycoris_locon_tab(headless=headless)
-            gradio_merge_lora_tab(headless=headless)
-            gradio_merge_lycoris_tab(headless=headless)
-            gradio_resize_lora_tab(headless=headless)
+            with gr.Tab('LoRA'):
+                gradio_extract_dylora_tab(headless=headless)
+                gradio_extract_lora_tab(headless=headless)
+                gradio_extract_lycoris_locon_tab(headless=headless)
+                gradio_merge_lora_tab(headless=headless)
+                gradio_merge_lycoris_tab(headless=headless)
+                gradio_resize_lora_tab(headless=headless)
 
     # Show the interface
     launch_kwargs = {}

library/group_images_gui.py

@@ -0,0 +1,110 @@
+import gradio as gr
+from easygui import msgbox
+import subprocess
+from .common_gui import get_folder_path
+import os
+
+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
+):
+    if input_folder == '':
+        msgbox('Input folder is missing...')
+        return
+
+    if output_folder == '':
+        msgbox('Please provide an output folder.')
+        return
+
+    print(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'
+
+    print(run_cmd)
+
+    if os.name == 'posix':
+        os.system(run_cmd)
+    else:
+        subprocess.run(run_cmd)
+
+    print('...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',
+            )
+
+        group_images_button = gr.Button('Caption images')
+
+        group_images_button.click(
+            group_images,
+            inputs=[
+                input_folder,
+                output_folder,
+                group_size,
+                include_subfolders,
+                do_not_copy_other_files
+            ],
+            show_progress=False,
+        )

library/utilities.py

@@ -11,6 +11,7 @@ from library.convert_model_gui import gradio_convert_model_tab
 from library.blip_caption_gui import gradio_blip_caption_gui_tab
 from library.git_caption_gui import gradio_git_caption_gui_tab
 from library.wd14_caption_gui import gradio_wd14_caption_gui_tab
+from library.group_images_gui import gradio_group_images_gui_tab
 
 
 def utilities_tab(
@@ -28,6 +29,7 @@ def utilities_tab(
         gradio_git_caption_gui_tab(headless=headless)
         gradio_wd14_caption_gui_tab(headless=headless)
     gradio_convert_model_tab(headless=headless)
+    gradio_group_images_gui_tab(headless=headless)
 
     return (
         train_data_dir_input,

tools/group_images.py

@@ -3,7 +3,7 @@ import shutil
 from PIL import Image
 import os
 import numpy as np
-import math
+# import math
 
 class ImageProcessor:
 
@@ -38,7 +38,7 @@ class ImageProcessor:
             cropped_images = self.crop_images(group, avg_aspect_ratio)
             self.resize_and_save_images(cropped_images, group_index)
             if not self.do_not_copy_other_files:
-                self.copy_other_files(group)
+                self.copy_other_files(group, group_index)
 
     def get_aspect_ratios(self, group):
         aspect_ratios = []
@@ -83,8 +83,8 @@ class ImageProcessor:
             print(f"  Saving processed image to {output_path}")
             img.convert('RGB').save(output_path)
 
-    def copy_other_files(self, group):
-        for path in group:
+    def copy_other_files(self, group, group_index):
+        for j, path in enumerate(group):
             dirpath, original_filename = os.path.split(path)
             original_basename, original_ext = os.path.splitext(original_filename)
             for filename in os.listdir(dirpath):
@@ -92,7 +92,7 @@ class ImageProcessor:
                     continue
                 basename, ext = os.path.splitext(filename)
                 if basename == original_basename and ext != original_ext:
-                    shutil.copy2(os.path.join(dirpath, filename), self.output_folder)
+                    shutil.copy2(os.path.join(dirpath, filename), os.path.join(self.output_folder, f"group-{group_index+1}-image-{j+1}{ext}"))
 
     def process_images(self):
         images = self.get_image_paths()