diff --git a/dataset_tag_editor.py b/dataset_tag_editor.py index 79fc42e..20befdb 100644 --- a/dataset_tag_editor.py +++ b/dataset_tag_editor.py @@ -17,6 +17,8 @@ class InterrogateMethod(Enum): NONE = 0 PREFILL = 1 OVERWRITE = 2 + PREPEND = 3 + APPEND = 4 def interrogate_image_clip(path): @@ -176,7 +178,7 @@ class DatasetTagEditor: - def load_dataset(self, img_dir: str, recursive: bool = False, load_caption_from_filename: bool = True, interrogate_method: InterrogateMethod = InterrogateMethod.NONE, use_clip: bool = True): + def load_dataset(self, img_dir: str, recursive: bool = False, load_caption_from_filename: bool = True, interrogate_method: InterrogateMethod = InterrogateMethod.NONE, use_booru: bool = True, use_clip: bool = False): self.clear() print(f'Loading dataset from {img_dir}') try: @@ -188,8 +190,8 @@ class DatasetTagEditor: self.dataset_dir = img_dir - if not (use_clip or cmd_opts.deepdanbooru): - print('Cannot interrogate without --deepdanbooru commandline option.') + if use_booru and not cmd_opts.deepdanbooru: + print('Cannot use DeepDanbooru without --deepdanbooru commandline option.') print(f'Total {len(filepath_set)} files under the directory including not image files.') @@ -214,16 +216,26 @@ class DatasetTagEditor: tokens = self.re_word.findall(caption_text) caption_text = (shared.opts.dataset_filename_join_string or "").join(tokens) - if interrogate_method == InterrogateMethod.OVERWRITE or (not caption_text and interrogate_method == InterrogateMethod.PREFILL): + if interrogate_method != InterrogateMethod.NONE and ((interrogate_method != InterrogateMethod.PREFILL) or (interrogate_method == InterrogateMethod.PREFILL and not caption_text)): try: img = Image.open(img_path).convert('RGB') except Exception as e: print(e) else: + interrogate_text = '' if use_clip: - caption_text = shared.interrogator.generate_caption(img) - elif cmd_opts.deepdanbooru: - caption_text = deepbooru.get_tags_from_process(img) + interrogate_text += shared.interrogator.generate_caption(img) + + if use_booru and cmd_opts.deepdanbooru: + tmp = deepbooru.get_tags_from_process(img) + interrogate_text += (', ' if interrogate_text and tmp else '') + tmp + + if interrogate_method == InterrogateMethod.OVERWRITE: + caption_text = interrogate_text + elif interrogate_method == InterrogateMethod.PREPEND: + caption_text = interrogate_text + (', ' if interrogate_text and caption_text else '') + caption_text + else: + caption_text += (', ' if interrogate_text and caption_text else '') + interrogate_text self.set_tags_by_image_path(img_path, [t.strip() for t in caption_text.split(',')]) @@ -231,7 +243,7 @@ class DatasetTagEditor: if interrogate_method != InterrogateMethod.NONE: if use_clip: shared.interrogator.load() - elif cmd_opts.deepdanbooru: + if use_booru and cmd_opts.deepdanbooru: db_opts = deepbooru.create_deepbooru_opts() db_opts[deepbooru.OPT_INCLUDE_RANKS] = False deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts) @@ -242,7 +254,7 @@ class DatasetTagEditor: if interrogate_method != InterrogateMethod.NONE: if use_clip: shared.interrogator.send_blip_to_ram() - elif cmd_opts.deepdanbooru: + if use_booru and cmd_opts.deepdanbooru: deepbooru.release_process() self.construct_tag_counts() diff --git a/scripts/dataset_tag_editor_ui.py b/scripts/dataset_tag_editor_ui.py index e125eea..6b2e9e0 100644 --- a/scripts/dataset_tag_editor_ui.py +++ b/scripts/dataset_tag_editor_ui.py @@ -38,7 +38,7 @@ def get_current_txt_selection(): return f"""Selected Image : {selection_selected_image_path}""" -def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool, load_caption_from_filename: bool, use_interrogator: str, use_clip: bool): +def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool, load_caption_from_filename: bool, use_interrogator: str, use_clip: bool, use_booru: bool): global total_image_num, displayed_image_num, current_tag_filter, current_selection, tmp_selection_img_path_set, selected_image_path, selection_selected_image_path interrogate_method = InterrogateMethod.NONE @@ -46,8 +46,12 @@ def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool interrogate_method = InterrogateMethod.PREFILL elif use_interrogator == 'Overwrite': interrogate_method = InterrogateMethod.OVERWRITE + elif use_interrogator == 'Prepend': + interrogate_method = InterrogateMethod.PREPEND + elif use_interrogator == 'Append': + interrogate_method = InterrogateMethod.APPEND - dataset_tag_editor.load_dataset(img_dir=dir, recursive=recursive, load_caption_from_filename=load_caption_from_filename, interrogate_method=interrogate_method, use_clip=use_clip) + dataset_tag_editor.load_dataset(img_dir=dir, recursive=recursive, load_caption_from_filename=load_caption_from_filename, interrogate_method=interrogate_method, use_clip=use_clip, use_booru=use_booru) img_paths, tags = dataset_tag_editor.get_filtered_imgpath_and_tags() tags = arrange_tag_order(tags=tags, sort_by=sort_by, sort_order=sort_order) total_image_num = displayed_image_num = len(dataset_tag_editor.get_img_path_set()) @@ -283,8 +287,10 @@ def on_ui_tabs(): cb_load_recursive = gr.Checkbox(value=False, label='Load from subdirectories') cb_load_caption_from_filename = gr.Checkbox(value=True, label='Load caption from filename if no text file') with gr.Column(): - rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite'], value='No', label='Use Interrogator Caption') - cb_use_clip_to_prefill = gr.Checkbox(value=False, label='Use BLIP interrogator') + rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite', 'Prepend', 'Append'], value='No', label='Use Interrogator Caption') + with gr.Row(): + cb_use_clip_to_prefill = gr.Checkbox(value=False, label='Use BLIP') + cb_use_booru_to_prefill = gr.Checkbox(value=False, label='Use DeepDanbooru') gl_dataset_images = gr.Gallery(label='Dataset Images', elem_id="dataset_tag_editor_dataset_gallery").style(grid=opts.dataset_editor_image_columns) txt_filter = gr.HTML(value=get_current_txt_filter()) @@ -381,7 +387,7 @@ def on_ui_tabs(): btn_load_datasets.click( fn=load_files_from_dir, - inputs=[tb_img_directory, rd_sort_by, rd_sort_order, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill], + inputs=[tb_img_directory, rd_sort_by, rd_sort_order, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill, cb_use_booru_to_prefill], outputs=[gl_dataset_images, gl_selected_images, cbg_tags, tb_search_tags, txt_filter, txt_selection] ) btn_load_datasets.click(