deepbooru: added option to use spaces or underscores

deepbooru: added option to quote (\) in tags
deepbooru/BLIP: write caption to file instead of image filename
deepbooru/BLIP: now possible to use both for captions
deepbooru: process is stopped even if an exception occurs

4 files changed, 95 insertions, 71 deletions

diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index 29529949..419e6a9c 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -2,33 +2,44 @@ import os.path
 from concurrent.futures import ProcessPoolExecutor
 import multiprocessing
 import time
+import re
+
+re_special = re.compile(r'([\\()])')
 
 def get_deepbooru_tags(pil_image):
     """
     This method is for running only one image at a time for simple use.  Used to the img2img interrogate.
     """
     from modules import shared  # prevents circular reference
-    create_deepbooru_process(shared.opts.interrogate_deepbooru_score_threshold, shared.opts.deepbooru_sort_alpha)
-    shared.deepbooru_process_return["value"] = -1
-    shared.deepbooru_process_queue.put(pil_image)
-    while shared.deepbooru_process_return["value"] == -1:
-        time.sleep(0.2)
-    tags = shared.deepbooru_process_return["value"]
-    release_process()
-    return tags
 
+    try:
+        create_deepbooru_process(shared.opts.interrogate_deepbooru_score_threshold, create_deepbooru_opts())
+        return get_tags_from_process(pil_image)
+    finally:
+        release_process()
+
+
+def create_deepbooru_opts():
+    from modules import shared
 
-def deepbooru_process(queue, deepbooru_process_return, threshold, alpha_sort):
+    return {
+        "use_spaces": shared.opts.deepbooru_use_spaces,
+        "use_escape": shared.opts.deepbooru_escape,
+        "alpha_sort": shared.opts.deepbooru_sort_alpha,
+    }
+
+
+def deepbooru_process(queue, deepbooru_process_return, threshold, deepbooru_opts):
     model, tags = get_deepbooru_tags_model()
     while True: # while process is running, keep monitoring queue for new image
         pil_image = queue.get()
         if pil_image == "QUIT":
             break
         else:
-            deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort)
+            deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold, deepbooru_opts)
 
 
-def create_deepbooru_process(threshold, alpha_sort):
+def create_deepbooru_process(threshold, deepbooru_opts):
     """
     Creates deepbooru process.  A queue is created to send images into the process.  This enables multiple images
     to be processed in a row without reloading the model or creating a new process.  To return the data, a shared
@@ -41,10 +52,23 @@ def create_deepbooru_process(threshold, alpha_sort):
     shared.deepbooru_process_queue = shared.deepbooru_process_manager.Queue()
     shared.deepbooru_process_return = shared.deepbooru_process_manager.dict()
     shared.deepbooru_process_return["value"] = -1
-    shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, alpha_sort))
+    shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, deepbooru_opts))
     shared.deepbooru_process.start()
 
 
+def get_tags_from_process(image):
+    from modules import shared
+
+    shared.deepbooru_process_return["value"] = -1
+    shared.deepbooru_process_queue.put(image)
+    while shared.deepbooru_process_return["value"] == -1:
+        time.sleep(0.2)
+    caption = shared.deepbooru_process_return["value"]
+    shared.deepbooru_process_return["value"] = -1
+
+    return caption
+
+
 def release_process():
     """
     Stops the deepbooru process to return used memory
@@ -81,10 +105,15 @@ def get_deepbooru_tags_model():
     return model, tags
 
 
-def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort):
+def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, deepbooru_opts):
     import deepdanbooru as dd
     import tensorflow as tf
     import numpy as np
+
+    alpha_sort = deepbooru_opts['alpha_sort']
+    use_spaces = deepbooru_opts['use_spaces']
+    use_escape = deepbooru_opts['use_escape']
+
     width = model.input_shape[2]
     height = model.input_shape[1]
     image = np.array(pil_image)
@@ -129,4 +158,12 @@ def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort)
 
     print('\n'.join(sorted(result_tags_print, reverse=True)))
 
-    return ', '.join(result_tags_out).replace('_', ' ').replace(':', ' ')
+    tags_text = ', '.join(result_tags_out)
+
+    if use_spaces:
+        tags_text = tags_text.replace('_', ' ')
+
+    if use_escape:
+        tags_text = re.sub(re_special, r'\\\1', tags_text)
+
+    return tags_text.replace(':', ' ')
diff --git a/modules/shared.py b/modules/shared.py
index e64e69fc..78b73aae 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -260,6 +260,8 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
     "interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),

     "interrogate_deepbooru_score_threshold": OptionInfo(0.5, "Interrogate: deepbooru score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),

     "deepbooru_sort_alpha": OptionInfo(True, "Interrogate: deepbooru sort alphabetically"),

+    "deepbooru_use_spaces": OptionInfo(False, "use spaces for tags in deepbooru"),

+    "deepbooru_escape": OptionInfo(True, "escape (\\) brackets in deepbooru (so they are used as literal brackets and not for emphasis)"),

 }))

 

 options_templates.update(options_section(('ui', "User interface"), {

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 113cecf1..3047bede 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -10,7 +10,28 @@ from modules.shared import opts, cmd_opts
 if cmd_opts.deepdanbooru:

     import modules.deepbooru as deepbooru

 

+

 def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

+    try:

+        if process_caption:

+            shared.interrogator.load()

+

+        if process_caption_deepbooru:

+            deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, deepbooru.create_deepbooru_opts())

+

+        preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru)

+

+    finally:

+

+        if process_caption:

+            shared.interrogator.send_blip_to_ram()

+

+        if process_caption_deepbooru:

+            deepbooru.release_process()

+

+

+

+def preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

     width = process_width

     height = process_height

     src = os.path.abspath(process_src)

@@ -25,30 +46,28 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
     shared.state.textinfo = "Preprocessing..."

     shared.state.job_count = len(files)

 

-    if process_caption:

-        shared.interrogator.load()

-

-    if process_caption_deepbooru:

-        deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, opts.deepbooru_sort_alpha)

-

     def save_pic_with_caption(image, index):

+        caption = ""

+

         if process_caption:

-            caption = "-" + shared.interrogator.generate_caption(image)

-            caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")

-        elif process_caption_deepbooru:

-            shared.deepbooru_process_return["value"] = -1

-            shared.deepbooru_process_queue.put(image)

-            while shared.deepbooru_process_return["value"] == -1:

-                time.sleep(0.2)

-            caption = "-" + shared.deepbooru_process_return["value"]

-            caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")

-            shared.deepbooru_process_return["value"] = -1

-        else:

-            caption = filename

-            caption = os.path.splitext(caption)[0]

-            caption = os.path.basename(caption)

+            caption += shared.interrogator.generate_caption(image)

+

+        if process_caption_deepbooru:

+            if len(caption) > 0:

+                caption += ", "

+            caption += deepbooru.get_tags_from_process(image)

+

+        filename_part = filename

+        filename_part = os.path.splitext(filename_part)[0]

+        filename_part = os.path.basename(filename_part)

+

+        basename = f"{index:05}-{subindex[0]}-{filename_part}"

+        image.save(os.path.join(dst, f"{basename}.png"))

+

+        if len(caption) > 0:

+            with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:

+                file.write(caption)

 

-        image.save(os.path.join(dst, f"{index:05}-{subindex[0]}{caption}.png"))

         subindex[0] += 1

 

     def save_pic(image, index):

@@ -93,34 +112,3 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
             save_pic(img, index)

 

         shared.state.nextjob()

-

-    if process_caption:

-        shared.interrogator.send_blip_to_ram()

-

-    if process_caption_deepbooru:

-        deepbooru.release_process()

-

-

-def sanitize_caption(base_path, original_caption, suffix):

-    operating_system = platform.system().lower()

-    if (operating_system == "windows"):

-        invalid_path_characters = "\\/:*?\"<>|"

-        max_path_length = 259

-    else:

-        invalid_path_characters = "/" #linux/macos

-        max_path_length = 1023

-    caption = original_caption

-    for invalid_character in invalid_path_characters:

-        caption = caption.replace(invalid_character, "")

-    fixed_path_length = len(base_path) + len(suffix) 

-    if fixed_path_length + len(caption) <= max_path_length:

-        return caption

-    caption_tokens = caption.split()

-    new_caption = ""

-    for token in caption_tokens:

-        last_caption = new_caption

-        new_caption = new_caption + token + " "

-        if (len(new_caption) + fixed_path_length - 1  > max_path_length):

-            break

-    print(f"\nPath will be too long. Truncated caption: {original_caption}\nto: {last_caption}", file=sys.stderr)

-    return last_caption.strip()

diff --git a/modules/ui.py b/modules/ui.py
index c42535c8..e07ee0e1 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1074,11 +1074,8 @@ def create_ui(wrap_gradio_gpu_call):
                     with gr.Row():

                         process_flip = gr.Checkbox(label='Create flipped copies')

                         process_split = gr.Checkbox(label='Split oversized images into two')

-                        process_caption = gr.Checkbox(label='Use BLIP caption as filename')

-                        if cmd_opts.deepdanbooru:

-                            process_caption_deepbooru = gr.Checkbox(label='Use deepbooru caption as filename')

-                        else:

-                            process_caption_deepbooru = gr.Checkbox(label='Use deepbooru caption as filename', visible=False)

+                        process_caption = gr.Checkbox(label='Use BLIP for caption')

+                        process_caption_deepbooru = gr.Checkbox(label='Use deepbooru for caption', visible=True if cmd_opts.deepdanbooru else False)

 

                     with gr.Row():

                         with gr.Column(scale=3):