Add new prefix and postfic for captions

library/basic_caption_gui.py

@@ -1,37 +1,40 @@
 import gradio as gr
 from easygui import msgbox
 import subprocess
-from .common_gui import get_folder_path
+from .common_gui import get_folder_path, add_pre_postfix
 
 
 def caption_images(
-    caption_text_input, images_dir_input, overwrite_input, caption_file_ext
+    caption_text_input, images_dir_input, overwrite_input, caption_file_ext, prefix, postfix
 ):
-    # Check for caption_text_input
-    if caption_text_input == '':
-        msgbox('Caption text is missing...')
-        return
-
     # Check for images_dir_input
     if images_dir_input == '':
         msgbox('Image folder is missing...')
         return
 
-    print(
+    if not caption_text_input == '':
-        f'Captioning files in {images_dir_input} with {caption_text_input}...'
+        print(
-    )
+            f'Captioning files in {images_dir_input} with {caption_text_input}...'
-    run_cmd = f'python "tools/caption.py"'
+        )
-    run_cmd += f' --caption_text="{caption_text_input}"'
+        run_cmd = f'python "tools/caption.py"'
+        run_cmd += f' --caption_text="{caption_text_input}"'
+        if overwrite_input:
+            run_cmd += f' --overwrite'
+        if caption_file_ext != '':
+            run_cmd += f' --caption_file_ext="{caption_file_ext}"'
+        run_cmd += f' "{images_dir_input}"'
+
+        print(run_cmd)
+
+        # Run the command
+        subprocess.run(run_cmd)
+    
     if overwrite_input:
-        run_cmd += f' --overwrite'
+        # Add prefix and postfix
-    if caption_file_ext != '':
+        add_pre_postfix(folder=images_dir_input, caption_file_ext=caption_file_ext, prefix=prefix, postfix=postfix)
-        run_cmd += f' --caption_file_ext="{caption_file_ext}"'
+    else:
-    run_cmd += f' "{images_dir_input}"'
+        if not prefix == '' or not postfix == '':
-
+            msgbox('Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected...')
-    print(run_cmd)
-
-    # Run the command
-    subprocess.run(run_cmd)
 
     print('...captioning done')
 
@@ -46,22 +49,6 @@ def gradio_basic_caption_gui_tab():
         gr.Markdown(
             'This utility will allow the creation of simple caption files for each images in a folder.'
         )
-        with gr.Row():
-            caption_text_input = gr.Textbox(
-                label='Caption text',
-                placeholder='Eg: , by some artist',
-                interactive=True,
-            )
-            overwrite_input = gr.Checkbox(
-                label='Overwrite existing captions in folder',
-                interactive=True,
-                value=False,
-            )
-            caption_file_ext = gr.Textbox(
-                label='Caption file extension',
-                placeholder='(Optional) Default: .caption',
-                interactive=True,
-            )
         with gr.Row():
             images_dir_input = gr.Textbox(
                 label='Image folder to caption',
@@ -74,6 +61,33 @@ def gradio_basic_caption_gui_tab():
             button_images_dir_input.click(
                 get_folder_path, outputs=images_dir_input
             )
+        with gr.Row():
+            prefix = gr.Textbox(
+                label='Prefix to add to txt caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+            caption_text_input = gr.Textbox(
+                label='Caption text',
+                placeholder='Eg: , by some artist. Leave empti if you just want to add pre or postfix',
+                interactive=True,
+            )
+            postfix = gr.Textbox(
+                label='Postfix to add to txt caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+        with gr.Row():
+            overwrite_input = gr.Checkbox(
+                label='Overwrite existing captions in folder',
+                interactive=True,
+                value=False,
+            )
+            caption_file_ext = gr.Textbox(
+                label='Caption file extension',
+                placeholder='(Optional) Default: .caption',
+                interactive=True,
+            )
         caption_button = gr.Button('Caption images')
 
         caption_button.click(
@@ -83,5 +97,6 @@ def gradio_basic_caption_gui_tab():
                 images_dir_input,
                 overwrite_input,
                 caption_file_ext,
+                prefix, postfix
             ],
         )

library/blip_caption_gui.py

@@ -1,8 +1,8 @@
 import gradio as gr
 from easygui import msgbox
 import subprocess
-from .common_gui import get_folder_path
+import os
-
+from .common_gui import get_folder_path, add_pre_postfix
 
 def caption_images(
     train_data_dir,
@@ -13,6 +13,8 @@ def caption_images(
     max_length,
     min_length,
     beam_search,
+    prefix,
+    postfix
 ):
     # Check for caption_text_input
     # if caption_text_input == "":
@@ -42,6 +44,9 @@ def caption_images(
 
     # Run the command
     subprocess.run(run_cmd)
+    
+    # Add prefix and postfix
+    add_pre_postfix(folder=train_data_dir, caption_file_ext=caption_file_ext, prefix=prefix, postfix=postfix)
 
     print('...captioning done')
 
@@ -68,12 +73,24 @@ def gradio_blip_caption_gui_tab():
             button_train_data_dir_input.click(
                 get_folder_path, outputs=train_data_dir
             )
-
+        with gr.Row():
             caption_file_ext = gr.Textbox(
                 label='Caption file extension',
                 placeholder='(Optional) Default: .caption',
                 interactive=True,
             )
+            
+            prefix = gr.Textbox(
+                label='Prefix to add to BLIP caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+            
+            postfix = gr.Textbox(
+                label='Postfix to add to BLIP caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
 
             batch_size = gr.Number(
                 value=1, label='Batch size', interactive=True
@@ -107,5 +124,7 @@ def gradio_blip_caption_gui_tab():
                 max_length,
                 min_length,
                 beam_search,
+                prefix,
+                postfix
             ],
         )

library/common_gui.py

@@ -1,4 +1,5 @@
 from tkinter import filedialog, Tk
+import os
 
 def get_file_path(file_path='', defaultextension='.json'):
     current_file_path = file_path
@@ -56,4 +57,19 @@ def get_saveasfile_path(file_path='', defaultextension='.json'):
 
     print(file_path)
 
-    return file_path
+    return file_path
+
+def add_pre_postfix(folder='', prefix='', postfix='', caption_file_ext='.caption'):
+    files  = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)]
+    if not prefix == '':
+        prefix = f'{prefix} '
+    if not postfix == '':
+        postfix = f' {postfix}'
+        
+    for file in files:
+        with open(os.path.join(folder, file), 'r+') as f:
+            content = f.read()
+            content = content.rstrip()
+            f.seek(0,0)
+            f.write(f'{prefix}{content}{postfix}')
+    f.close()