Code refactor

library/basic_caption_gui.py

@@ -20,43 +20,51 @@ def caption_images(
     find_text,
     replace_text,
 ):
-    # Check for images_dir
+    # Check if images_dir is provided
     if not images_dir:
-        msgbox('Image folder is missing...')
+        msgbox('Image folder is missing. Please provide the directory containing the images to caption.')
         return
 
+    # Check if caption_ext is provided
     if not caption_ext:
         msgbox('Please provide an extension for the caption files.')
         return
 
     if caption_text:
         log.info(f'Captioning files in {images_dir} with {caption_text}...')
+
+        # Build the command to run caption.py
         run_cmd = f'python "tools/caption.py"'
         run_cmd += f' --caption_text="{caption_text}"'
+
+        # Add optional flags to the command
         if overwrite:
             run_cmd += f' --overwrite'
         if caption_ext:
             run_cmd += f' --caption_file_ext="{caption_ext}"'
+
         run_cmd += f' "{images_dir}"'
 
         log.info(run_cmd)
 
-        # Run the command
+        # Run the command based on the operating system
         if os.name == 'posix':
             os.system(run_cmd)
         else:
             subprocess.run(run_cmd)
 
+    # Check if overwrite option is enabled
     if overwrite:
         if prefix or postfix:
-            # Add prefix and postfix
+            # Add prefix and postfix to caption files
             add_pre_postfix(
                 folder=images_dir,
                 caption_file_ext=caption_ext,
                 prefix=prefix,
                 postfix=postfix,
             )
         if find_text:
+            # Find and replace text in caption files
             find_replace(
                 folder_path=images_dir,
                 caption_file_ext=caption_ext,
@@ -65,18 +73,19 @@ def caption_images(
             )
     else:
         if prefix or postfix:
+            # Show a message if modification is not possible without overwrite option enabled
             msgbox(
-                'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected...'
+                'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.'
             )
 
-    log.info('...captioning done')
+    log.info('Captioning done.')
 
 
 # Gradio UI
 def gradio_basic_caption_gui_tab(headless=False):
     with gr.Tab('Basic Captioning'):
         gr.Markdown(
-            'This utility will allow the creation of simple caption files for each image in a folder.'
+            'This utility allows you to create simple caption files for each image in a folder.'
         )
         with gr.Row():
             images_dir = gr.Textbox(
@@ -94,7 +103,7 @@ def gradio_basic_caption_gui_tab(headless=False):
             )
             caption_ext = gr.Textbox(
                 label='Caption file extension',
-                placeholder='Extension for caption file. eg: .caption, .txt',
+                placeholder='Extension for caption file (e.g., .caption, .txt)',
                 value='.txt',
                 interactive=True,
             )
@@ -111,7 +120,7 @@ def gradio_basic_caption_gui_tab(headless=False):
             )
             caption_text = gr.Textbox(
                 label='Caption text',
-                placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix',
+                placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix or postfix.',
                 interactive=True,
             )
             postfix = gr.Textbox(
@@ -122,12 +131,12 @@ def gradio_basic_caption_gui_tab(headless=False):
         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',
+                placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix 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',
+                placeholder='e.g., "by some artist". Leave empty if you want to replace with nothing.',
                 interactive=True,
             )
             caption_button = gr.Button('Caption images')

library/blip_caption_gui.py

@@ -3,7 +3,6 @@
 import subprocess
 import os
 from .common_gui import get_folder_path, add_pre_postfix
-
 from library.custom_logging import setup_logging
 
 # Set up logging
@@ -24,21 +23,19 @@ def caption_images(
     prefix,
     postfix,
 ):
-    # Check for caption_text_input
-    # if caption_text_input == "":
-    #     msgbox("Caption text is missing...")
-    #     return
-
-    # Check for images_dir_input
+    # Check if the image folder is provided
     if train_data_dir == '':
         msgbox('Image folder is missing...')
         return
 
+    # Check if the caption file extension is provided
     if caption_file_ext == '':
         msgbox('Please provide an extension for the caption files.')
         return
 
     log.info(f'Captioning files in {train_data_dir}...')
+
+    # Construct the command to run
     run_cmd = f'{PYTHON} "finetune/make_captions.py"'
     run_cmd += f' --batch_size="{int(batch_size)}"'
     run_cmd += f' --num_beams="{int(num_beams)}"'
@@ -79,7 +76,7 @@ def caption_images(
 def gradio_blip_caption_gui_tab(headless=False):
     with gr.Tab('BLIP Captioning'):
         gr.Markdown(
-            'This utility will use BLIP to caption files for each images in a folder.'
+            'This utility uses BLIP to caption files for each image in a folder.'
         )
         with gr.Row():
             train_data_dir = gr.Textbox(
@@ -98,7 +95,7 @@ def gradio_blip_caption_gui_tab(headless=False):
         with gr.Row():
             caption_file_ext = gr.Textbox(
                 label='Caption file extension',
-                placeholder='Extention for caption file. eg: .caption, .txt',
+                placeholder='Extension for caption file, e.g., .caption, .txt',
                 value='.txt',
                 interactive=True,
             )

tools/group_images.py

@@ -92,18 +92,19 @@ def resize_and_save_images(self, cropped_images, group_index, source_paths):
             os.makedirs(self.output_folder, exist_ok=True)
             original_filename = os.path.basename(source_paths[j])
             filename_without_ext = os.path.splitext(original_filename)[0]
-            output_path = os.path.join(self.output_folder, f"group-{group_index+1}-{j+1}-{filename_without_ext}.jpg")
+            final_file_name = f"group-{group_index+1}-{j+1}-{filename_without_ext}"
+            output_path = os.path.join(self.output_folder, f"{final_file_name}.jpg")
             log.info(f"  Saving processed image to {output_path}")
             img.convert('RGB').save(output_path, quality=100)
 
             if self.caption:
-                self.create_caption_file(source_paths[j], group_index, filename_without_ext)
+                self.create_caption_file(source_paths[j], group_index, final_file_name)
 
     def create_caption_file(self, source_path, group_index, caption_filename):
         dirpath = os.path.dirname(source_path)
         caption = os.path.basename(dirpath).split('_')[-1]
         caption_filename = caption_filename + self.caption_ext
-        caption_path = os.path.join(self.output_folder, f"group-{group_index+1}-{caption_filename}")
+        caption_path = os.path.join(self.output_folder, caption_filename)
         with open(caption_path, 'w') as f:
             f.write(caption)