diff --git a/README-JP.md b/README-JP.md index 1c2ee18..026eb0d 100644 --- a/README-JP.md +++ b/README-JP.md @@ -16,8 +16,7 @@ DeepBooru interrogator で生成したような、カンマ区切り形式のキ ### WebUIのExtensionsタブからインストールする "Install from URL" タブに `https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git` をコピーしてインストールできます。 "Availables" タブにこの拡張機能が表示されている場合は、ワンクリックでインストール可能です。 -~~web UI の "Extensions" タブから更新をした際、完全に更新を適用するには web UI を再起動する必要がある場合があります。~~ -web UIの "Extensions" タブからの更新で問題なくリロードできるようになりました。 +**web UI の "Extensions" タブから更新をした際、完全に更新を適用するには web UI を再起動する必要がある場合があります。** ### 手動でインストールする web UI の `extensions` フォルダにリポジトリのクローンを作成し再起動してください。 diff --git a/README.md b/README.md index 251f31d..8f8742b 100644 --- a/README.md +++ b/README.md @@ -16,8 +16,7 @@ Copy `https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor. Also, if you see this extension listed, you can install from "Available" tab with a single click. -~~Please note that if you update this extension from "Extensions" tab, you will need to restart web UI to reload completely.~~ -Now you can reload by pushing "Apply and restart UI" button on Extensions tab! +**Please note that if you update this extension from "Extensions" tab, you will need to restart web UI to reload completely.** ### Install Manually To install, clone the repository into the `extensions` directory and restart the web UI. diff --git a/scripts/dataset_tag_editor/dataset_tag_editor.py b/scripts/dataset_tag_editor/__init__.py similarity index 93% rename from scripts/dataset_tag_editor/dataset_tag_editor.py rename to scripts/dataset_tag_editor/__init__.py index c03fd0c..9094670 100644 --- a/scripts/dataset_tag_editor/dataset_tag_editor.py +++ b/scripts/dataset_tag_editor/__init__.py @@ -6,14 +6,15 @@ from modules.textual_inversion.dataset import re_numbers_at_start from PIL import Image from enum import Enum -from scripts.dynamic_import import dynamic_import -ds = dynamic_import('scripts/dataset_tag_editor/dataset.py') -tagger = dynamic_import('scripts/dataset_tag_editor/tagger.py') -captioning = dynamic_import('scripts/dataset_tag_editor/captioning.py') -filters = dynamic_import('scripts/dataset_tag_editor/filters.py') -kohya_metadata = dynamic_import('scripts/dataset_tag_editor/kohya-ss_finetune_metadata.py') +from . import dataset as ds +from . import tagger +from . import captioning +from . import filters +from . import kohya_finetune_metadata as kohya_metadata +__all__ = ["ds", "tagger", "captioning", "filters", "kohya_metadata", "INTERROGATOR_NAMES", "InterrogateMethod", "interrogate_image", "DatasetTagEditor"] + re_tags = re.compile(r'^(.+) \[\d+\]$') WD_TAGGER_NAMES = ["wd-v1-4-vit-tagger", "wd-v1-4-convnext-tagger", "wd-v1-4-vit-tagger-v2", "wd-v1-4-convnext-tagger-v2", "wd-v1-4-swinv2-tagger-v2"] @@ -22,6 +23,7 @@ WD_TAGGER_THRESHOLDS = [0.35, 0.35, 0.3537, 0.3685, 0.3771] # v1: idk if it's ok INTERROGATORS = [captioning.BLIP(), tagger.DeepDanbooru()] + [tagger.WaifuDiffusion(name, WD_TAGGER_THRESHOLDS[i]) for i, name in enumerate(WD_TAGGER_NAMES)] INTERROGATOR_NAMES = [it.name() for it in INTERROGATORS] + class InterrogateMethod(Enum): NONE = 0 PREFILL = 1 @@ -155,35 +157,39 @@ class DatasetTagEditor: tags:Set[str] = self.dataset.get_tagset() result = set() - for tag in tags: - if prefix: - if regex: - if re.search("^" + filter_word, tag) is not None: - result.add(tag) - continue - else: - if tag.startswith(filter_word): - result.add(tag) - continue - if suffix: - if regex: - if re.search(filter_word + "$", tag) is not None: - result.add(tag) - continue - else: - if tag.endswith(filter_word): - result.add(tag) - continue - if not prefix and not suffix: - if regex: - if re.search(filter_word, tag) is not None: - result.add(tag) - continue - else: - if filter_word in tag: - result.add(tag) - continue - return result + try: + for tag in tags: + if prefix: + if regex: + if re.search("^" + filter_word, tag) is not None: + result.add(tag) + continue + else: + if tag.startswith(filter_word): + result.add(tag) + continue + if suffix: + if regex: + if re.search(filter_word + "$", tag) is not None: + result.add(tag) + continue + else: + if tag.endswith(filter_word): + result.add(tag) + continue + if not prefix and not suffix: + if regex: + if re.search(filter_word, tag) is not None: + result.add(tag) + continue + else: + if filter_word in tag: + result.add(tag) + continue + except: + return tags + else: + return result def cleanup_tags(self, tags: List[str]): @@ -400,7 +406,7 @@ class DatasetTagEditor: if move_image: try: dst_path_obj = dest_dir_obj / img_path_obj.name - if dst_path_obj.is_file(): + if img_path_obj.is_file(): if img_path in self.images: self.images[img_path].close() del self.images[img_path] diff --git a/scripts/dataset_tag_editor/captioning.py b/scripts/dataset_tag_editor/captioning.py index f4bd405..044a843 100644 --- a/scripts/dataset_tag_editor/captioning.py +++ b/scripts/dataset_tag_editor/captioning.py @@ -1,8 +1,8 @@ import modules.shared as shared -from scripts.dataset_tag_editor.interrogator import Interrogator -from scripts.dynamic_import import dynamic_import -git_large_captioning = dynamic_import('scripts/dataset_tag_editor/interrogators/git_large_captioning.py') +from .interrogator import Interrogator +from .interrogators import GITLargeCaptioning + class Captioning(Interrogator): def start(self): @@ -10,9 +10,9 @@ class Captioning(Interrogator): def stop(self): pass def predict(self, image): - raise NotImplementedError + raise NotImplementedError() def name(self): - raise NotImplementedError + raise NotImplementedError() class BLIP(Captioning): @@ -31,14 +31,16 @@ class BLIP(Captioning): class GITLarge(Captioning): + def __init__(self): + self.interrogator = GITLargeCaptioning() def start(self): - git_large_captioning.instance.load() + self.interrogator.load() def stop(self): - git_large_captioning.instance.unload() + self.interrogator.unload() def predict(self, image): - tags = git_large_captioning.instance.apply(image).split(',') + tags = self.interrogator.apply(image).split(',') return [t for t in tags if t] def name(self): diff --git a/scripts/dataset_tag_editor/interrogator.py b/scripts/dataset_tag_editor/interrogator.py index d5a0a8e..38408d4 100644 --- a/scripts/dataset_tag_editor/interrogator.py +++ b/scripts/dataset_tag_editor/interrogator.py @@ -10,6 +10,6 @@ class Interrogator: def stop(self): pass def predict(self, image, **kwargs): - raise NotImplementedError + raise NotImplementedError() def name(self): - raise NotImplementedError + raise NotImplementedError() diff --git a/scripts/dataset_tag_editor/interrogators/__init__.py b/scripts/dataset_tag_editor/interrogators/__init__.py new file mode 100644 index 0000000..726c896 --- /dev/null +++ b/scripts/dataset_tag_editor/interrogators/__init__.py @@ -0,0 +1,6 @@ +from .git_large_captioning import GITLargeCaptioning +from .waifu_diffusion_tagger import WaifuDiffusionTagger + +__all__ = [ + 'GITLargeCaptioning', 'WaifuDiffusionTagger' +] \ No newline at end of file diff --git a/scripts/dataset_tag_editor/interrogators/git_large_captioning.py b/scripts/dataset_tag_editor/interrogators/git_large_captioning.py index ddf5490..2a46767 100644 --- a/scripts/dataset_tag_editor/interrogators/git_large_captioning.py +++ b/scripts/dataset_tag_editor/interrogators/git_large_captioning.py @@ -23,7 +23,4 @@ class GITLargeCaptioning(): return '' inputs = self.processor(images=image, return_tensors='pt').to(shared.device) ids = self.model.generate(pixel_values=inputs.pixel_values, max_length=shared.opts.interrogate_clip_max_length) - return self.processor.batch_decode(ids, skip_special_tokens=True)[0] - - -instance = GITLargeCaptioning() \ No newline at end of file + return self.processor.batch_decode(ids, skip_special_tokens=True)[0] \ No newline at end of file diff --git a/scripts/dataset_tag_editor/kohya-ss_finetune_metadata.py b/scripts/dataset_tag_editor/kohya_finetune_metadata.py similarity index 100% rename from scripts/dataset_tag_editor/kohya-ss_finetune_metadata.py rename to scripts/dataset_tag_editor/kohya_finetune_metadata.py diff --git a/scripts/dataset_tag_editor/tagger.py b/scripts/dataset_tag_editor/tagger.py index 42a2e30..5ee520b 100644 --- a/scripts/dataset_tag_editor/tagger.py +++ b/scripts/dataset_tag_editor/tagger.py @@ -6,9 +6,8 @@ from typing import Optional, Dict from modules import devices, shared from modules import deepbooru as db -from scripts.dataset_tag_editor.interrogator import Interrogator -from scripts.dynamic_import import dynamic_import -waifu_diffusion_tagger = dynamic_import('scripts/dataset_tag_editor/interrogators/waifu_diffusion_tagger.py') +from .interrogator import Interrogator +from .interrogators import WaifuDiffusionTagger class Tagger(Interrogator): @@ -17,9 +16,9 @@ class Tagger(Interrogator): def stop(self): pass def predict(self, image: Image.Image, threshold: Optional[float]): - raise NotImplementedError + raise NotImplementedError() def name(self): - raise NotImplementedError + raise NotImplementedError() def get_replaced_tag(tag: str): @@ -76,7 +75,7 @@ class DeepDanbooru(Tagger): class WaifuDiffusion(Tagger): def __init__(self, repo_name, threshold): self.repo_name = repo_name - self.tagger_inst = waifu_diffusion_tagger.WaifuDiffusionTagger("SmilingWolf/" + repo_name) + self.tagger_inst = WaifuDiffusionTagger("SmilingWolf/" + repo_name) self.threshold = threshold def start(self): diff --git a/scripts/dataset_tag_editor/ui.py b/scripts/dataset_tag_editor/ui.py deleted file mode 100644 index 4e702e5..0000000 --- a/scripts/dataset_tag_editor/ui.py +++ /dev/null @@ -1,233 +0,0 @@ -from typing import List, Callable - -from scripts.dynamic_import import dynamic_import -dte = dynamic_import('scripts/dataset_tag_editor/dataset_tag_editor.py') -filters = dte.filters - - -class TagFilterUI: - def __init__(self, dataset_tag_editor, tag_filter_mode = filters.TagFilter.Mode.INCLUSIVE): - self.logic = filters.TagFilter.Logic.AND - self.filter_word = '' - self.sort_by = 'Alphabetical Order' - self.sort_order = 'Ascending' - self.selected_tags = set() - self.filter_mode = tag_filter_mode - self.filter = filters.TagFilter(logic=self.logic, mode=self.filter_mode) - self.dataset_tag_editor = dataset_tag_editor - self.get_filters = lambda:[] - self.prefix = False - self.suffix = False - self.regex = False - - def get_filter(self): - return self.filter - - def create_ui(self, get_filters: Callable[[], List[filters.Filter]], logic = filters.TagFilter.Logic.AND, sort_by = 'Alphabetical Order', sort_order = 'Ascending', prefix=False, suffix=False, regex=False): - self.get_filters = get_filters - self.logic = logic - self.filter = filters.TagFilter(logic=self.logic, mode=self.filter_mode) - self.sort_by = sort_by - self.sort_order = sort_order - self.prefix = prefix - self.suffix = suffix - self.regex = regex - - import gradio as gr - self.tb_search_tags = gr.Textbox(label='Search Tags', interactive=True) - with gr.Row(): - self.cb_prefix = gr.Checkbox(label='Prefix', value=self.prefix, interactive=True) - self.cb_suffix = gr.Checkbox(label='Suffix', value=self.suffix, interactive=True) - self.cb_regex = gr.Checkbox(label='Use regex', value=self.regex, interactive=True) - with gr.Row(): - self.rb_sort_by = gr.Radio(choices=['Alphabetical Order', 'Frequency', 'Length'], value=sort_by, interactive=True, label='Sort by') - self.rb_sort_order = gr.Radio(choices=['Ascending', 'Descending'], value=sort_order, interactive=True, label='Sort Order') - v = 'AND' if self.logic==filters.TagFilter.Logic.AND else 'OR' if self.logic==filters.TagFilter.Logic.OR else 'NONE' - self.rb_logic = gr.Radio(choices=['AND', 'OR', 'NONE'], value=v, label='Filter Logic', interactive=True) - self.cbg_tags = gr.CheckboxGroup(label='Filter Images by Tags', interactive=True) - - - def set_callbacks(self, on_filter_update: Callable[[List], List] = lambda:[], inputs=[], outputs=[], _js=None): - self.tb_search_tags.change(fn=self.tb_search_tags_changed, inputs=self.tb_search_tags, outputs=self.cbg_tags) - self.cb_prefix.change(fn=self.cb_prefix_changed, inputs=self.cb_prefix, outputs=self.cbg_tags) - self.cb_suffix.change(fn=self.cb_suffix_changed, inputs=self.cb_suffix, outputs=self.cbg_tags) - self.cb_regex.change(fn=self.cb_regex_changed, inputs=self.cb_regex, outputs=self.cbg_tags) - self.rb_sort_by.change(fn=self.rd_sort_by_changed, inputs=self.rb_sort_by, outputs=self.cbg_tags) - self.rb_sort_order.change(fn=self.rd_sort_order_changed, inputs=self.rb_sort_order, outputs=self.cbg_tags) - self.rb_logic.change(fn=lambda a, *b:[self.rd_logic_changed(a)] + on_filter_update(*b), _js=_js, inputs=[self.rb_logic] + inputs, outputs=[self.cbg_tags] + outputs) - self.cbg_tags.change(fn=lambda a, *b:[self.cbg_tags_changed(a)] + on_filter_update(*b), _js=_js, inputs=[self.cbg_tags] + inputs, outputs=[self.cbg_tags] + outputs) - - - - def tb_search_tags_changed(self, tb_search_tags: str): - self.filter_word = tb_search_tags - return self.cbg_tags_update() - - - def cb_prefix_changed(self, prefix:bool): - self.prefix = prefix - return self.cbg_tags_update() - - - def cb_suffix_changed(self, suffix:bool): - self.suffix = suffix - return self.cbg_tags_update() - - - def cb_regex_changed(self, use_regex:bool): - self.regex = use_regex - return self.cbg_tags_update() - - - def rd_sort_by_changed(self, rb_sort_by: str): - self.sort_by = rb_sort_by - return self.cbg_tags_update() - - - def rd_sort_order_changed(self, rd_sort_order: str): - self.sort_order = rd_sort_order - return self.cbg_tags_update() - - - def rd_logic_changed(self, rd_logic: str): - self.logic = filters.TagFilter.Logic.AND if rd_logic == 'AND' else filters.TagFilter.Logic.OR if rd_logic == 'OR' else filters.TagFilter.Logic.NONE - self.filter = filters.TagFilter(self.selected_tags, self.logic, self.filter_mode) - return self.cbg_tags_update() - - - def cbg_tags_changed(self, cbg_tags: List[str]): - self.selected_tags = self.dataset_tag_editor.cleanup_tagset(set(self.dataset_tag_editor.read_tags(cbg_tags))) - return self.cbg_tags_update() - - - def cbg_tags_update(self): - self.selected_tags = self.dataset_tag_editor.cleanup_tagset(self.selected_tags) - self.filter = filters.TagFilter(self.selected_tags, self.logic, self.filter_mode) - - if self.filter_mode == filters.TagFilter.Mode.INCLUSIVE: - tags = self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, self.filter.logic == filters.TagFilter.Logic.AND, prefix=self.prefix, suffix=self.suffix, regex=self.regex) - else: - tags = self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, self.filter.logic == filters.TagFilter.Logic.OR, prefix=self.prefix, suffix=self.suffix, regex=self.regex) - tags_in_filter = self.filter.tags - - tags = self.dataset_tag_editor.sort_tags(tags=tags, sort_by=self.sort_by, sort_order=self.sort_order) - tags_in_filter = self.dataset_tag_editor.sort_tags(tags=tags_in_filter, sort_by=self.sort_by, sort_order=self.sort_order) - - tags = tags_in_filter + [tag for tag in tags if tag not in self.filter.tags] - tags = self.dataset_tag_editor.write_tags(tags) - tags_in_filter = self.dataset_tag_editor.write_tags(tags_in_filter) - - import gradio as gr - return gr.CheckboxGroup.update(value=tags_in_filter, choices=tags) - - - def clear_filter(self): - self.filter = filters.TagFilter(logic=self.logic, mode=self.filter_mode) - self.filter_word = '' - self.selected_tags = set() - - -class TagSelectUI: - def __init__(self, dataset_tag_editor): - self.filter_word = '' - self.sort_by = 'Alphabetical Order' - self.sort_order = 'Ascending' - self.selected_tags = set() - self.tags = set() - self.dataset_tag_editor = dataset_tag_editor - self.get_filters = lambda:[] - self.prefix = False - self.suffix = False - self.regex = False - - - def create_ui(self, get_filters: Callable[[], List[filters.Filter]], sort_by = 'Alphabetical Order', sort_order = 'Ascending', prefix=False, suffix=False, regex=False): - self.get_filters = get_filters - self.prefix = prefix - self.suffix = suffix - self.regex = regex - - import gradio as gr - self.tb_search_tags = gr.Textbox(label='Search Tags', interactive=True) - with gr.Row(): - self.cb_prefix = gr.Checkbox(label='Prefix', value=False, interactive=True) - self.cb_suffix = gr.Checkbox(label='Suffix', value=False, interactive=True) - self.cb_regex = gr.Checkbox(label='Use regex', value=False, interactive=True) - with gr.Row(): - self.rb_sort_by = gr.Radio(choices=['Alphabetical Order', 'Frequency', 'Length'], value=sort_by, interactive=True, label='Sort by') - self.rb_sort_order = gr.Radio(choices=['Ascending', 'Descending'], value=sort_order, interactive=True, label='Sort Order') - with gr.Row(): - self.btn_select_visibles = gr.Button(value='Select visible tags') - self.btn_deselect_visibles = gr.Button(value='Deselect visible tags') - self.cbg_tags = gr.CheckboxGroup(label='Select Tags', interactive=True) - - - def set_callbacks(self): - self.tb_search_tags.change(fn=self.tb_search_tags_changed, inputs=self.tb_search_tags, outputs=self.cbg_tags) - self.cb_prefix.change(fn=self.cb_prefix_changed, inputs=self.cb_prefix, outputs=self.cbg_tags) - self.cb_suffix.change(fn=self.cb_suffix_changed, inputs=self.cb_suffix, outputs=self.cbg_tags) - self.cb_regex.change(fn=self.cb_regex_changed, inputs=self.cb_regex, outputs=self.cbg_tags) - self.rb_sort_by.change(fn=self.rd_sort_by_changed, inputs=self.rb_sort_by, outputs=self.cbg_tags) - self.rb_sort_order.change(fn=self.rd_sort_order_changed, inputs=self.rb_sort_order, outputs=self.cbg_tags) - self.btn_select_visibles.click(fn=self.btn_select_visibles_clicked, outputs=self.cbg_tags) - self.btn_deselect_visibles.click(fn=self.btn_deselect_visibles_clicked, inputs=self.cbg_tags, outputs=self.cbg_tags) - self.cbg_tags.change(fn=self.cbg_tags_changed, inputs=self.cbg_tags, outputs=self.cbg_tags) - - - def tb_search_tags_changed(self, tb_search_tags: str): - self.filter_word = tb_search_tags - return self.cbg_tags_update() - - - def cb_prefix_changed(self, prefix:bool): - self.prefix = prefix - return self.cbg_tags_update() - - - def cb_suffix_changed(self, suffix:bool): - self.suffix = suffix - return self.cbg_tags_update() - - - def cb_regex_changed(self, regex:bool): - self.regex = regex - return self.cbg_tags_update() - - - def rd_sort_by_changed(self, rb_sort_by: str): - self.sort_by = rb_sort_by - return self.cbg_tags_update() - - - def rd_sort_order_changed(self, rd_sort_order: str): - self.sort_order = rd_sort_order - return self.cbg_tags_update() - - - def cbg_tags_changed(self, cbg_tags: List[str]): - self.selected_tags = set(self.dataset_tag_editor.read_tags(cbg_tags)) - return self.cbg_tags_update() - - - def btn_deselect_visibles_clicked(self, cbg_tags: List[str]): - tags = self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, True) - selected_tags = set(self.dataset_tag_editor.read_tags(cbg_tags)) & tags - self.selected_tags -= selected_tags - return self.cbg_tags_update() - - - def btn_select_visibles_clicked(self): - tags = set(self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, True)) - self.selected_tags |= tags - return self.cbg_tags_update() - - - def cbg_tags_update(self): - tags = self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, True, prefix=self.prefix, suffix=self.suffix, regex=self.regex) - self.tags = set(self.dataset_tag_editor.get_filtered_tags(self.get_filters(), filter_tags=True, prefix=self.prefix, suffix=self.suffix, regex=self.regex)) - self.selected_tags &= self.tags - tags = self.dataset_tag_editor.sort_tags(tags=tags, sort_by=self.sort_by, sort_order=self.sort_order) - tags = self.dataset_tag_editor.write_tags(tags) - selected_tags = self.dataset_tag_editor.write_tags(list(self.selected_tags)) - import gradio as gr - return gr.CheckboxGroup.update(value=selected_tags, choices=tags) \ No newline at end of file diff --git a/scripts/dte_instance.py b/scripts/dte_instance.py new file mode 100644 index 0000000..a174e79 --- /dev/null +++ b/scripts/dte_instance.py @@ -0,0 +1,2 @@ +import scripts.dataset_tag_editor as dte_module +dte_instance = dte_module.DatasetTagEditor() \ No newline at end of file diff --git a/scripts/dynamic_import.py b/scripts/dynamic_import.py deleted file mode 100644 index 4a26d5a..0000000 --- a/scripts/dynamic_import.py +++ /dev/null @@ -1,5 +0,0 @@ -def dynamic_import(path: str): - import os - from modules import scripts, script_loading - path = os.path.abspath(os.path.join(scripts.basedir(), path)) - return script_loading.load_module(path) \ No newline at end of file diff --git a/scripts/main.py b/scripts/main.py index 0024441..6600ebd 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -6,28 +6,7 @@ import json from pathlib import Path from collections import namedtuple -from scripts.dynamic_import import dynamic_import -ui = dynamic_import('scripts/dataset_tag_editor/ui.py') -filters = ui.filters -dte = ui.dte - - -dataset_tag_editor = dte.DatasetTagEditor() - -path_filter = filters.PathFilter() -tag_filter_ui:ui.TagFilterUI = None -tag_filter_ui_neg:ui.TagFilterUI = None -tag_select_ui_remove:ui.TagSelectUI = None - -def get_filters(): - return [path_filter, tag_filter_ui.get_filter(), tag_filter_ui_neg.get_filter()] - -total_image_num = 0 -displayed_image_num = 0 -tmp_selection_img_path_set = set() -gallery_selected_image_path = '' -selection_selected_image_path = '' -show_only_selected_tags = True +import scripts.ui as ui # ================================================================ @@ -172,357 +151,36 @@ def read_move_delete_config(): return read_config('file_move_delete', MoveDeleteConfig, CFG_MOVE_DELETE_DEFAULT) # ================================================================ -# Callbacks for "Filter and Edit Tags" tab +# General Callbacks for Updating UIs # ================================================================ -def get_current_gallery_txt(): - return f""" - Displayed Images : {displayed_image_num} / {total_image_num} total
- Current Tag Filter : {tag_filter_ui.get_filter()} {' AND ' if tag_filter_ui.get_filter().tags and tag_filter_ui_neg.get_filter().tags else ''} {tag_filter_ui_neg.get_filter()}
- Current Selection Filter : {len(path_filter.paths)} images
- Selected Image : {gallery_selected_image_path} - """ - - -def get_current_txt_selection(): - return f"""Selected Image : {selection_selected_image_path}""" - - -def get_current_move_or_delete_target_num(target_data: str, idx: int): - if target_data == 'Selected One': - idx = int(idx) - return f'Target dataset num: {1 if idx != -1 else 0}' - - elif target_data == 'All Displayed Ones': - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - return f'Target dataset num: {len(img_paths)}' - - else: - return f'Target dataset num: 0' - - -def load_files_from_dir( - dir: str, - caption_file_ext: str, - recursive: bool, - load_caption_from_filename: bool, - use_interrogator: str, - use_interrogator_names: List[str], - use_custom_threshold_booru: bool, - custom_threshold_booru: float, - use_custom_threshold_waifu: bool, - custom_threshold_waifu: float, - use_kohya_metadata: bool, - kohya_json_path: str - ): - global total_image_num, displayed_image_num, tmp_selection_img_path_set, gallery_selected_image_path, selection_selected_image_path, path_filter - - interrogate_method = dte.InterrogateMethod.NONE - if use_interrogator == 'If Empty': - interrogate_method = dte.InterrogateMethod.PREFILL - elif use_interrogator == 'Overwrite': - interrogate_method = dte.InterrogateMethod.OVERWRITE - elif use_interrogator == 'Prepend': - interrogate_method = dte.InterrogateMethod.PREPEND - elif use_interrogator == 'Append': - interrogate_method = dte.InterrogateMethod.APPEND - - threshold_booru = custom_threshold_booru if use_custom_threshold_booru else shared.opts.interrogate_deepbooru_score_threshold - threshold_waifu = custom_threshold_waifu if use_custom_threshold_waifu else -1 - - dataset_tag_editor.load_dataset(dir, caption_file_ext, recursive, load_caption_from_filename, interrogate_method, use_interrogator_names, threshold_booru, threshold_waifu, opts.dataset_editor_use_temp_files, kohya_json_path if use_kohya_metadata else None) - imgs = dataset_tag_editor.get_filtered_imgs(filters=[]) - img_indices = dataset_tag_editor.get_filtered_imgindices(filters=[]) - path_filter = filters.PathFilter() - total_image_num = displayed_image_num = len(dataset_tag_editor.get_img_path_set()) - tmp_selection_img_path_set = set() - gallery_selected_image_path = '' - selection_selected_image_path = '' - return [ - imgs, - [], - get_current_gallery_txt(), - get_current_txt_selection() - ] +\ - [gr.CheckboxGroup.update(value=[str(i) for i in img_indices], choices=[str(i) for i in img_indices]), True] +\ - clear_tag_filters() +\ - [tag_select_ui_remove.cbg_tags_update()] - +def get_filters(): + filters = [ui.filter_by_tags.tag_filter_ui.get_filter(), ui.filter_by_tags.tag_filter_ui_neg.get_filter()] + [ui.filter_by_selection.path_filter] + return filters def update_gallery(): - global displayed_image_num - img_indices = dataset_tag_editor.get_filtered_imgindices(filters=get_filters()) + img_indices = ui.dte_instance.get_filtered_imgindices(filters=get_filters()) + total_image_num = len(ui.dte_instance.dataset) displayed_image_num = len(img_indices) + ui.gallery_state.register_value('Displayed Images', f'{displayed_image_num} / {total_image_num} total') + ui.gallery_state.register_value('Current Tag Filter', f"{ui.filter_by_tags.tag_filter_ui.get_filter()} {' AND ' if ui.filter_by_tags.tag_filter_ui.get_filter().tags and ui.filter_by_tags.tag_filter_ui_neg.get_filter().tags else ''} {ui.filter_by_tags.tag_filter_ui_neg.get_filter()}") + ui.gallery_state.register_value('Current Selection Filter', f'{len(ui.filter_by_selection.path_filter.paths)} images') return [ [str(i) for i in img_indices], 1, -1, -1, -1, - get_current_gallery_txt() + ui.gallery_state.get_current_gallery_txt() ] - def update_filter_and_gallery(): return \ - [tag_filter_ui.cbg_tags_update(), tag_filter_ui_neg.cbg_tags_update()] +\ + [ui.filter_by_tags.tag_filter_ui.cbg_tags_update(), ui.filter_by_tags.tag_filter_ui_neg.cbg_tags_update()] +\ update_gallery() +\ - update_common_tags() +\ - [', '.join(tag_filter_ui.filter.tags)] +\ - [tag_select_ui_remove.cbg_tags_update()] - - -def clear_tag_filters(): - tag_filter_ui.clear_filter() - tag_filter_ui_neg.clear_filter() - return update_filter_and_gallery() - - -def update_common_tags(): - if show_only_selected_tags: - tags = ', '.join([t for t in dataset_tag_editor.get_common_tags(filters=get_filters()) if t in tag_filter_ui.filter.tags]) - else: - tags = ', '.join(dataset_tag_editor.get_common_tags(filters=get_filters())) - return [tags, tags] - - -def save_all_changes(backup: bool, caption_ext: str, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool): - if not metadata_input: - metadata_input = None - saved, total, dir = dataset_tag_editor.save_dataset(backup, caption_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath) - return f'Saved text files : {saved}/{total} under {dir}' if total > 0 else '' - - -# ================================================================ -# Callbacks for "Filter by Selection" tab -# ================================================================ - -def arrange_selection_order(paths: List[str]): - return sorted(paths) - - -def selection_index_changed(idx: int): - global tmp_selection_img_path_set, selection_selected_image_path - idx = int(idx) - img_paths = arrange_selection_order(tmp_selection_img_path_set) - if idx < 0 or len(img_paths) <= idx: - selection_selected_image_path = '' - idx = -1 - else: - selection_selected_image_path = img_paths[idx] - return [get_current_txt_selection(), idx] - - -def add_image_selection(idx: int): - global tmp_selection_img_path_set - idx = int(idx) - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - if idx < 0 or len(img_paths) <= idx: - idx = -1 - else: - tmp_selection_img_path_set.add(img_paths[idx]) - return [arrange_selection_order(tmp_selection_img_path_set), idx] - - -def add_all_displayed_image_selection(): - global tmp_selection_img_path_set - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - tmp_selection_img_path_set |= set(img_paths) - return arrange_selection_order(tmp_selection_img_path_set) - - -def invert_image_selection(): - global tmp_selection_img_path_set - img_paths = dataset_tag_editor.get_img_path_set() - tmp_selection_img_path_set = img_paths - tmp_selection_img_path_set - return arrange_selection_order(tmp_selection_img_path_set) - - -def remove_image_selection(idx: int): - global tmp_selection_img_path_set, selection_selected_image_path - idx = int(idx) - img_paths = arrange_selection_order(tmp_selection_img_path_set) - if 0 < idx and idx < len(img_paths): - tmp_selection_img_path_set.remove(img_paths[idx]) - selection_selected_image_path = '' - idx = -1 - - return [ - arrange_selection_order(tmp_selection_img_path_set), - get_current_txt_selection(), - idx - ] - - -def clear_image_selection(): - global tmp_selection_img_path_set, selection_selected_image_path, path_filter - tmp_selection_img_path_set.clear() - selection_selected_image_path = '' - path_filter = filters.PathFilter() - return[ - [], - get_current_txt_selection(), - -1 - ] - - -def apply_image_selection_filter(): - global path_filter - if len(tmp_selection_img_path_set) > 0: - path_filter = filters.PathFilter(tmp_selection_img_path_set, filters.PathFilter.Mode.INCLUSIVE) - else: - path_filter = filters.PathFilter() - return update_filter_and_gallery() - - -# ================================================================ -# Callbacks for "Edit Caption of Selected Image" tab -# ================================================================ - -def gallery_index_changed(prev_idx: int, next_idx: int, edit_caption: str, copy_automatically: bool, warn_change_not_saved: bool): - global gallery_selected_image_path - prev_idx = int(prev_idx) - next_idx = int(next_idx) - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - prev_tags_txt = '' - if 0 <= prev_idx and prev_idx < len(img_paths): - prev_tags_txt = ', '.join(dataset_tag_editor.get_tags_by_image_path(gallery_selected_image_path)) - else: - prev_idx = -1 - - next_tags_txt = '' - if 0 <= next_idx and next_idx < len(img_paths): - gallery_selected_image_path = img_paths[next_idx] - next_tags_txt = ', '.join(dataset_tag_editor.get_tags_by_image_path(gallery_selected_image_path)) - else: - gallery_selected_image_path = '' - - return\ - [prev_idx if warn_change_not_saved and edit_caption != prev_tags_txt else -1] +\ - [next_tags_txt, next_tags_txt if copy_automatically else edit_caption] +\ - [edit_caption] +\ - [get_current_gallery_txt()] - - -def dialog_selected_save_caption_change(prev_idx: int, edit_caption: str, sort: bool = False): - prev_idx = int(prev_idx) - return change_selected_image_caption(edit_caption, prev_idx, sort) - - -def change_selected_image_caption(tags_text: str, idx: int, sort: bool = False): - idx = int(idx) - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - - edited_tags = [t.strip() for t in tags_text.split(',')] - edited_tags = [t for t in edited_tags if t] - - if sort: - edited_tags = dataset_tag_editor.sort_tags(edited_tags) - - if 0 <= idx and idx < len(img_paths): - dataset_tag_editor.set_tags_by_image_path(imgpath=img_paths[idx], tags=edited_tags) - return update_filter_and_gallery() - - -def interrogate_selected_image(interrogator_name: str, use_threshold_booru: bool, threshold_booru: float, use_threshold_waifu: bool, threshold_waifu: float): - global gallery_selected_image_path - threshold_booru = threshold_booru if use_threshold_booru else shared.opts.interrogate_deepbooru_score_threshold - threshold_waifu = threshold_waifu if use_threshold_waifu else -1 - return dte.interrogate_image(gallery_selected_image_path, interrogator_name, threshold_booru, threshold_waifu) - - -# ================================================================ -# Callbacks for "Batch Edit Captions" tab -# ================================================================ - -def apply_edit_tags(search_tags: str, replace_tags: str, prepend: bool): - search_tags = [t.strip() for t in search_tags.split(',')] - search_tags = [t for t in search_tags if t] - replace_tags = [t.strip() for t in replace_tags.split(',')] - replace_tags = [t for t in replace_tags if t] - - dataset_tag_editor.replace_tags(search_tags = search_tags, replace_tags = replace_tags, filters=get_filters(), prepend = prepend) - tag_filter_ui.get_filter().tags = dataset_tag_editor.get_replaced_tagset(tag_filter_ui.get_filter().tags, search_tags, replace_tags) - tag_filter_ui_neg.get_filter().tags = dataset_tag_editor.get_replaced_tagset(tag_filter_ui_neg.get_filter().tags, search_tags, replace_tags) - - return update_filter_and_gallery() - - -def search_and_replace(search_text: str, replace_text: str, target_text: str, use_regex: bool): - if target_text == 'Only Selected Tags': - selected_tags = set(tag_filter_ui.selected_tags) - dataset_tag_editor.search_and_replace_selected_tags(search_text = search_text, replace_text=replace_text, selected_tags=selected_tags, filters=get_filters(), use_regex=use_regex) - tag_filter_ui.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui.filter.tags, selected_tags, use_regex) - tag_filter_ui_neg.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui_neg.filter.tags, selected_tags, use_regex) - - elif target_text == 'Each Tags': - dataset_tag_editor.search_and_replace_selected_tags(search_text = search_text, replace_text=replace_text, selected_tags=None, filters=get_filters(), use_regex=use_regex) - tag_filter_ui.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui.filter.tags, None, use_regex) - tag_filter_ui_neg.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui_neg.filter.tags, None, use_regex) - - elif target_text == 'Entire Caption': - dataset_tag_editor.search_and_replace_caption(search_text=search_text, replace_text=replace_text, filters=get_filters(), use_regex=use_regex) - tag_filter_ui.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui.filter.tags, None, use_regex) - tag_filter_ui_neg.filter.tags = dataset_tag_editor.search_and_replace_tag_set(search_text, replace_text, tag_filter_ui_neg.filter.tags, None, use_regex) - - return update_filter_and_gallery() - - -def cb_show_only_tags_selected_changed(value: bool): - global show_only_selected_tags - show_only_selected_tags = value - return update_common_tags() - - -def remove_duplicated_tags(): - dataset_tag_editor.remove_duplicated_tags(get_filters()) - return update_filter_and_gallery() - - -def remove_selected_tags(): - dataset_tag_editor.remove_tags(tag_select_ui_remove.selected_tags, get_filters()) - return update_filter_and_gallery() - -def sort_selected_tags(**sort_args): - dataset_tag_editor.sort_filtered_tags(get_filters(), **sort_args) - return update_filter_and_gallery() - -# ================================================================ -# Callbacks for "Move or Delete Files" tab -# ================================================================ - -def move_files(target_data: str, target_file: List[str], caption_ext: str, dest_dir: str, idx: int = -1): - move_img = 'Image File' in target_file - move_txt = 'Caption Text File' in target_file - move_bak = 'Caption Backup File' in target_file - if target_data == 'Selected One': - idx = int(idx) - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - if 0 <= idx and idx < len(img_paths): - dataset_tag_editor.move_dataset_file(img_paths[idx], caption_ext, dest_dir, move_img, move_txt, move_bak) - dataset_tag_editor.construct_tag_counts() - - elif target_data == 'All Displayed Ones': - dataset_tag_editor.move_dataset(dest_dir, caption_ext, get_filters(), move_img, move_txt, move_bak) - - return update_filter_and_gallery() - - -def delete_files(target_data: str, target_file: List[str], caption_ext: str, idx: int = -1): - delete_img = 'Image File' in target_file - delete_txt = 'Caption Text File' in target_file - delete_bak = 'Caption Backup File' in target_file - if target_data == 'Selected One': - idx = int(idx) - img_paths = dataset_tag_editor.get_filtered_imgpaths(filters=get_filters()) - if 0 <= idx and idx < len(img_paths): - dataset_tag_editor.delete_dataset_file(delete_img, caption_ext, delete_txt, delete_bak) - dataset_tag_editor.construct_tag_counts() - - elif target_data == 'All Displayed Ones': - dataset_tag_editor.delete_dataset(caption_ext, get_filters(), delete_img, delete_txt, delete_bak) - - return update_filter_and_gallery() + ui.batch_edit_captions.get_common_tags(get_filters, ui.filter_by_tags) +\ + [', '.join(ui.filter_by_tags.tag_filter_ui.filter.tags)] +\ + [ui.batch_edit_captions.tag_select_ui_remove.cbg_tags_update()] # ================================================================ @@ -530,7 +188,6 @@ def delete_files(target_data: str, target_file: List[str], caption_ext: str, idx # ================================================================ def on_ui_tabs(): - global tag_filter_ui, tag_filter_ui_neg, tag_select_ui_remove config.load() cfg_general = read_general_config() @@ -539,44 +196,14 @@ def on_ui_tabs(): cfg_edit_selected = read_edit_selected_config() cfg_file_move_delete = read_move_delete_config() - tag_filter_ui = ui.TagFilterUI(dataset_tag_editor, tag_filter_mode=filters.TagFilter.Mode.INCLUSIVE) - tag_filter_ui_neg = ui.TagFilterUI(dataset_tag_editor, tag_filter_mode=filters.TagFilter.Mode.EXCLUSIVE) - tag_select_ui_remove = ui.TagSelectUI(dataset_tag_editor) - with gr.Blocks(analytics_enabled=False) as dataset_tag_editor_interface: - with gr.Row(visible=False): - btn_hidden_set_index = gr.Button(elem_id="dataset_tag_editor_btn_hidden_set_index") - nb_hidden_image_index = gr.Number(value=-1, label='hidden_idx_next') - nb_hidden_image_index_prev = gr.Number(value=-1, label='hidden_idx_prev') - nb_hidden_image_index_save_or_not = gr.Number(value=-1, label='hidden_s_or_n') - btn_hidden_save_caption = gr.Button(elem_id="dataset_tag_editor_btn_hidden_save_caption") - tb_hidden_edit_caption = gr.Textbox() - cbg_hidden_dataset_filter = gr.CheckboxGroup(label='Dataset Filter') - nb_hidden_dataset_filter_apply = gr.Number(label='Filter Apply', value=-1) gr.HTML(value=""" This extension works well with text captions in comma-separated style (such as the tags generated by DeepBooru interrogator). """) - with gr.Column(variant='panel'): - with gr.Row(): - with gr.Column(scale=1): - btn_save_all_changes = gr.Button(value='Save all changes', variant='primary') - with gr.Column(scale=2): - cb_backup = gr.Checkbox(value=cfg_general.backup, label='Backup original text file (original file will be renamed like filename.000, .001, .002, ...)', interactive=True) - gr.HTML(value='Note: New text file will be created if you are using filename as captions.') - with gr.Row(): - cb_save_kohya_metadata = gr.Checkbox(value=cfg_general.save_kohya_metadata, label="Use kohya-ss's finetuning metadata json", interactive=True) - with gr.Row(): - with gr.Column(variant='panel', visible=cfg_general.save_kohya_metadata) as kohya_metadata: - tb_metadata_output = gr.Textbox(label='json path', placeholder='C:\\path\\to\\metadata.json',value=cfg_general.meta_output_path) - tb_metadata_input = gr.Textbox(label='json input path (Optional, only for append results)', placeholder='C:\\path\\to\\metadata.json',value=cfg_general.meta_input_path) - with gr.Row(): - cb_metadata_overwrite = gr.Checkbox(value=cfg_general.meta_overwrite, label="Overwrite if output file exists", interactive=True) - cb_metadata_as_caption = gr.Checkbox(value=cfg_general.meta_save_as_caption, label="Save metadata as caption", interactive=True) - cb_metadata_use_fullpath = gr.Checkbox(value=cfg_general.meta_use_full_path, label="Save metadata image key as fullpath", interactive=True) - with gr.Row(visible=False): - txt_result = gr.Textbox(label='Results', interactive=False) - + + ui.toprow.create_ui(cfg_general) + with gr.Accordion(label='Reload/Save Settings (config.json)', open=False): with gr.Row(): btn_reload_config_file = gr.Button(value='Reload settings') @@ -585,176 +212,40 @@ def on_ui_tabs(): with gr.Row().style(equal_height=False): with gr.Column(): - with gr.Column(variant='panel'): - with gr.Row(): - with gr.Column(scale=3): - tb_img_directory = gr.Textbox(label='Dataset directory', placeholder='C:\\directory\\of\\datasets', value=cfg_general.dataset_dir) - with gr.Column(scale=1, min_width=60): - tb_caption_file_ext = gr.Textbox(label='Caption File Ext', placeholder='.txt (on Load and Save)', value=cfg_general.caption_ext) - with gr.Column(scale=1, min_width=80): - btn_load_datasets = gr.Button(value='Load') - with gr.Accordion(label='Dataset Load Settings'): - with gr.Row(): - with gr.Column(): - cb_load_recursive = gr.Checkbox(value=cfg_general.load_recursive, label='Load from subdirectories') - cb_load_caption_from_filename = gr.Checkbox(value=cfg_general.load_caption_from_filename, label='Load caption from filename if no text file exists') - with gr.Column(): - rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite', 'Prepend', 'Append'], value=cfg_general.use_interrogator, label='Use Interrogator Caption') - dd_intterogator_names = gr.Dropdown(label = 'Interrogators', choices=dte.INTERROGATOR_NAMES, value=cfg_general.use_interrogator_names, interactive=True, multiselect=True) - with gr.Accordion(label='Interrogator Settings', open=False): - with gr.Row(): - cb_use_custom_threshold_booru = gr.Checkbox(value=cfg_general.use_custom_threshold_booru, label='Use Custom Threshold (Booru)', interactive=True) - sl_custom_threshold_booru = gr.Slider(minimum=0, maximum=1, value=cfg_general.custom_threshold_booru, step=0.01, interactive=True, label='Booru Score Threshold') - with gr.Row(): - cb_use_custom_threshold_waifu = gr.Checkbox(value=cfg_general.use_custom_threshold_waifu, label='Use Custom Threshold (WDv1.4 Tagger)', interactive=True) - sl_custom_threshold_waifu = gr.Slider(minimum=0, maximum=1, value=cfg_general.custom_threshold_waifu, step=0.01, interactive=True, label='WDv1.4 Tagger Score Threshold') - - gl_dataset_images = gr.Gallery(label='Dataset Images', elem_id="dataset_tag_editor_dataset_gallery").style(grid=opts.dataset_editor_image_columns) - txt_gallery = gr.HTML(value=get_current_gallery_txt()) + ui.load_dataset.create_ui(cfg_general) + ui.dataset_gallery.create_ui(opts.dataset_editor_image_columns) + ui.gallery_state.create_ui() with gr.Tab(label='Filter by Tags'): - with gr.Row(): - btn_clear_tag_filters = gr.Button(value='Clear tag filters') - btn_clear_all_filters = gr.Button(value='Clear ALL filters') - - with gr.Tab(label='Positive Filter'): - with gr.Column(variant='panel'): - gr.HTML(value='Search tags / Filter images by tags (INCLUSIVE)') - logic_p = filters.TagFilter.Logic.OR if cfg_filter_p.logic=='OR' else filters.TagFilter.Logic.NONE if cfg_filter_p.logic=='NONE' else filters.TagFilter.Logic.AND - tag_filter_ui.create_ui(get_filters, logic_p, cfg_filter_p.sort_by, cfg_filter_p.sort_order, cfg_filter_p.sw_prefix, cfg_filter_p.sw_suffix, cfg_filter_p.sw_regex) - - with gr.Tab(label='Negative Filter'): - with gr.Column(variant='panel'): - gr.HTML(value='Search tags / Filter images by tags (EXCLUSIVE)') - logic_n = filters.TagFilter.Logic.AND if cfg_filter_n.logic=='AND' else filters.TagFilter.Logic.NONE if cfg_filter_n.logic=='NONE' else filters.TagFilter.Logic.OR - tag_filter_ui_neg.create_ui(get_filters, logic_n, cfg_filter_n.sort_by, cfg_filter_n.sort_order, cfg_filter_n.sw_prefix, cfg_filter_n.sw_suffix, cfg_filter_n.sw_regex) + ui.filter_by_tags.create_ui(cfg_filter_p, cfg_filter_n, get_filters) with gr.Tab(label='Filter by Selection'): - with gr.Row(visible=False): - btn_hidden_set_selection_index = gr.Button(elem_id="dataset_tag_editor_btn_hidden_set_selection_index") - nb_hidden_selection_image_index = gr.Number(value=-1) - gr.HTML("""Select images from the left gallery.""") - - with gr.Column(variant='panel'): - with gr.Row(): - btn_add_image_selection = gr.Button(value='Add selection [Enter]', elem_id='dataset_tag_editor_btn_add_image_selection') - btn_add_all_displayed_image_selection = gr.Button(value='Add ALL Displayed') - - gl_filter_images = gr.Gallery(label='Filter Images', elem_id="dataset_tag_editor_filter_gallery").style(grid=opts.dataset_editor_image_columns) - txt_selection = gr.HTML(value=get_current_txt_selection()) - - with gr.Row(): - btn_remove_image_selection = gr.Button(value='Remove selection [Delete]', elem_id='dataset_tag_editor_btn_remove_image_selection') - btn_invert_image_selection = gr.Button(value='Invert selection') - btn_clear_image_selection = gr.Button(value='Clear selection') - - btn_apply_image_selection_filter = gr.Button(value='Apply selection filter', variant='primary') + ui.filter_by_selection.create_ui(opts.dataset_editor_image_columns) with gr.Tab(label='Batch Edit Captions'): - with gr.Tab(label='Search and Replace'): - with gr.Column(variant='panel'): - gr.HTML('Edit common tags.') - cb_show_only_tags_selected = gr.Checkbox(value=cfg_batch_edit.show_only_selected, label='Show only the tags selected in the Positive Filter') - tb_common_tags = gr.Textbox(label='Common Tags', interactive=False) - tb_edit_tags = gr.Textbox(label='Edit Tags', interactive=True) - cb_prepend_tags = gr.Checkbox(value=cfg_batch_edit.prepend, label='Prepend additional tags') - btn_apply_edit_tags = gr.Button(value='Apply changes to filtered images', variant='primary') - with gr.Accordion(label='Show description of how to edit tags', open=False): - gr.HTML(value=""" - 1. The tags common to all displayed images are shown in comma separated style.
- 2. When changes are applied, all tags in each displayed images are replaced.
- 3. If you change some tags into blank, they will be erased.
- 4. If you add some tags to the end, they will be added to the end/beginning of the text file.
- 5. Changes are not applied to the text files until the "Save all changes" button is pressed.
- ex A.
-  Original Text = "A, A, B, C" Common Tags = "B, A" Edit Tags = "X, Y"
-  Result = "Y, Y, X, C" (B->X, A->Y)
- ex B.
-  Original Text = "A, B, C" Common Tags = "(nothing)" Edit Tags = "X, Y"
-  Result = "A, B, C, X, Y" (add X and Y to the end (default))
-  Result = "X, Y, A, B, C" (add X and Y to the beginning ("Prepend additional tags" checked))
- ex C.
-  Original Text = "A, B, C, D, E" Common Tags = "A, B, D" Edit Tags = ", X, "
-  Result = "X, C, E" (A->"", B->X, D->"")
- """) - with gr.Column(variant='panel'): - gr.HTML('Search and Replace for all images displayed.') - tb_sr_search_tags = gr.Textbox(label='Search Text', interactive=True) - tb_sr_replace_tags = gr.Textbox(label='Replace Text', interactive=True) - cb_use_regex = gr.Checkbox(label='Use regex', value=cfg_batch_edit.use_regex) - rb_sr_replace_target = gr.Radio(['Only Selected Tags', 'Each Tags', 'Entire Caption'], value=cfg_batch_edit.target, label='Search and Replace in', interactive=True) - tb_sr_selected_tags = gr.Textbox(label='Selected Tags', interactive=False, lines=2) - btn_apply_sr_tags = gr.Button(value='Search and Replace', variant='primary') - with gr.Tab(label='Remove'): - with gr.Column(variant='panel'): - gr.HTML('Remove duplicate tags from the images displayed.') - btn_remove_duplicate = gr.Button(value='Remove duplicate tags', variant='primary') - with gr.Column(variant='panel'): - gr.HTML('Remove selected tags from the images displayed.') - btn_remove_selected = gr.Button(value='Remove selected tags', variant='primary') - tag_select_ui_remove.create_ui(get_filters, cfg_batch_edit.sory_by, cfg_batch_edit.sort_order, cfg_batch_edit.sw_prefix, cfg_batch_edit.sw_suffix, cfg_batch_edit.sw_regex) - with gr.Tab(label='Extras'): - with gr.Column(variant='panel'): - gr.HTML('Sort tags from the images displayed.') - btn_sort_selected = gr.Button(value='Sort selected tags', variant='primary') + ui.batch_edit_captions.create_ui(cfg_batch_edit, get_filters) with gr.Tab(label='Edit Caption of Selected Image'): - with gr.Tab(label='Read Caption from Selected Image'): - tb_caption_selected_image = gr.Textbox(label='Caption of Selected Image', interactive=True, lines=6) - with gr.Row(): - btn_copy_caption = gr.Button(value='Copy and Overwrite') - btn_prepend_caption = gr.Button(value='Prepend') - btn_append_caption = gr.Button(value='Append') - - with gr.Tab(label='Interrogate Selected Image'): - with gr.Row(): - dd_intterogator_names_si = gr.Dropdown(label = 'Interrogator', choices=dte.INTERROGATOR_NAMES, value=cfg_edit_selected.use_interrogator_name, interactive=True, multiselect=False) - btn_interrogate_si = gr.Button(value='Interrogate') - tb_interrogate_selected_image = gr.Textbox(label='Interrogate Result', interactive=True, lines=6) - with gr.Row(): - btn_copy_interrogate = gr.Button(value='Copy and Overwrite') - btn_prepend_interrogate = gr.Button(value='Prepend') - btn_append_interrogate = gr.Button(value='Append') - cb_copy_caption_automatically = gr.Checkbox(value=cfg_edit_selected.auto_copy, label='Copy caption from selected images automatically') - cb_sort_caption_on_save = gr.Checkbox(value=cfg_edit_selected.sort_on_save, label='Sort caption on save') - cb_ask_save_when_caption_changed = gr.Checkbox(value=cfg_edit_selected.warn_change_not_saved, label='Warn if changes in caption is not saved') - tb_edit_caption_selected_image = gr.Textbox(label='Edit Caption', interactive=True, lines=6) - btn_apply_changes_selected_image = gr.Button(value='Apply changes to selected image', variant='primary') - - gr.HTML("""Changes are not applied to the text files until the "Save all changes" button is pressed.""") + ui.edit_caption_of_selected_image.create_ui(cfg_edit_selected) with gr.Tab(label='Move or Delete Files'): - gr.HTML(value='Note: Moved or deleted images will be unloaded.') - rb_move_or_delete_target_data = gr.Radio(choices=['Selected One', 'All Displayed Ones'], value=cfg_file_move_delete.range, label='Move or Delete') - cbg_move_or_delete_target_file = gr.CheckboxGroup(choices=['Image File', 'Caption Text File', 'Caption Backup File'], label='Target', value=cfg_file_move_delete.target) - tb_move_or_delete_caption_ext = gr.Textbox(label='Caption File Ext', placeholder='txt', value=cfg_file_move_delete.caption_ext) - ta_move_or_delete_target_dataset_num = gr.HTML(value='Target dataset num: 0') - tb_move_or_delete_destination_dir = gr.Textbox(label='Destination Directory', value=cfg_file_move_delete.destination) - btn_move_or_delete_move_files = gr.Button(value='Move File(s)', variant='primary') - gr.HTML(value='Note: DELETE cannot be undone. The files will be deleted completely.') - btn_move_or_delete_delete_files = gr.Button(value='DELETE File(s)', variant='primary') - - o_filter_and_gallery = \ - [tag_filter_ui.cbg_tags, tag_filter_ui_neg.cbg_tags, cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply, nb_hidden_image_index, nb_hidden_image_index_prev, nb_hidden_image_index_save_or_not, txt_gallery] +\ - [tb_common_tags, tb_edit_tags] +\ - [tb_sr_selected_tags] +\ - [tag_select_ui_remove.cbg_tags] + ui.move_or_delete_files.create_ui(cfg_file_move_delete) #---------------------------------------------------------------- # General components_general = [ - cb_backup, tb_img_directory, tb_caption_file_ext, cb_load_recursive, - cb_load_caption_from_filename, rb_use_interrogator, dd_intterogator_names, - cb_use_custom_threshold_booru, sl_custom_threshold_booru, cb_use_custom_threshold_waifu, sl_custom_threshold_waifu, - cb_save_kohya_metadata, tb_metadata_output, tb_metadata_input, cb_metadata_overwrite, cb_metadata_as_caption, cb_metadata_use_fullpath + ui.toprow.cb_backup, ui.load_dataset.tb_img_directory, ui.load_dataset.tb_caption_file_ext, ui.load_dataset.cb_load_recursive, + ui.load_dataset.cb_load_caption_from_filename, ui.load_dataset.rb_use_interrogator, ui.load_dataset.dd_intterogator_names, + ui.load_dataset.cb_use_custom_threshold_booru, ui.load_dataset.sl_custom_threshold_booru, ui.load_dataset.cb_use_custom_threshold_waifu, ui.load_dataset.sl_custom_threshold_waifu, + ui.toprow.cb_save_kohya_metadata, ui.toprow.tb_metadata_output, ui.toprow.tb_metadata_input, ui.toprow.cb_metadata_overwrite, ui.toprow.cb_metadata_as_caption, ui.toprow.cb_metadata_use_fullpath ] components_filter = \ - [tag_filter_ui.cb_prefix, tag_filter_ui.cb_suffix, tag_filter_ui.cb_regex, tag_filter_ui.rb_sort_by, tag_filter_ui.rb_sort_order, tag_filter_ui.rb_logic] +\ - [tag_filter_ui_neg.cb_prefix, tag_filter_ui_neg.cb_suffix, tag_filter_ui_neg.cb_regex, tag_filter_ui_neg.rb_sort_by, tag_filter_ui_neg.rb_sort_order, tag_filter_ui_neg.rb_logic] - components_batch_edit = [cb_show_only_tags_selected, cb_prepend_tags, cb_use_regex, rb_sr_replace_target, tag_select_ui_remove.cb_prefix, tag_select_ui_remove.cb_suffix, tag_select_ui_remove.cb_regex, tag_select_ui_remove.rb_sort_by, tag_select_ui_remove.rb_sort_order] - components_edit_selected = [cb_copy_caption_automatically, cb_sort_caption_on_save, cb_ask_save_when_caption_changed, dd_intterogator_names_si] - components_move_delete = [rb_move_or_delete_target_data, cbg_move_or_delete_target_file, tb_move_or_delete_caption_ext, tb_move_or_delete_destination_dir] + [ui.filter_by_tags.tag_filter_ui.cb_prefix, ui.filter_by_tags.tag_filter_ui.cb_suffix, ui.filter_by_tags.tag_filter_ui.cb_regex, ui.filter_by_tags.tag_filter_ui.rb_sort_by, ui.filter_by_tags.tag_filter_ui.rb_sort_order, ui.filter_by_tags.tag_filter_ui.rb_logic] +\ + [ui.filter_by_tags.tag_filter_ui_neg.cb_prefix, ui.filter_by_tags.tag_filter_ui_neg.cb_suffix, ui.filter_by_tags.tag_filter_ui_neg.cb_regex, ui.filter_by_tags.tag_filter_ui_neg.rb_sort_by, ui.filter_by_tags.tag_filter_ui_neg.rb_sort_order, ui.filter_by_tags.tag_filter_ui_neg.rb_logic] + components_batch_edit = [ui.batch_edit_captions.cb_show_only_tags_selected, ui.batch_edit_captions.cb_prepend_tags, ui.batch_edit_captions.cb_use_regex, ui.batch_edit_captions.rb_sr_replace_target, ui.batch_edit_captions.tag_select_ui_remove.cb_prefix, ui.batch_edit_captions.tag_select_ui_remove.cb_suffix, ui.batch_edit_captions.tag_select_ui_remove.cb_regex, ui.batch_edit_captions.tag_select_ui_remove.rb_sort_by, ui.batch_edit_captions.tag_select_ui_remove.rb_sort_order] + components_edit_selected = [ui.edit_caption_of_selected_image.cb_copy_caption_automatically, ui.edit_caption_of_selected_image.cb_sort_caption_on_save, ui.edit_caption_of_selected_image.cb_ask_save_when_caption_changed, ui.edit_caption_of_selected_image.dd_intterogator_names_si] + components_move_delete = [ui.move_or_delete_files.rb_move_or_delete_target_data, ui.move_or_delete_files.cbg_move_or_delete_target_file, ui.move_or_delete_files.tb_move_or_delete_caption_ext, ui.move_or_delete_files.tb_move_or_delete_destination_dir] configurable_components = components_general + components_filter + components_batch_edit + components_edit_selected + components_move_delete @@ -803,335 +294,25 @@ def on_ui_tabs(): outputs=configurable_components ) - btn_save_all_changes.click( - fn=save_all_changes, - inputs=[cb_backup, tb_caption_file_ext, cb_save_kohya_metadata, tb_metadata_output, tb_metadata_input, cb_metadata_overwrite, cb_metadata_as_caption, cb_metadata_use_fullpath], - outputs=[txt_result] - ) + o_update_gallery = [ui.filter_by_tags.cbg_hidden_dataset_filter, ui.filter_by_tags.nb_hidden_dataset_filter_apply, ui.dataset_gallery.nb_hidden_image_index, ui.dataset_gallery.nb_hidden_image_index_prev, ui.edit_caption_of_selected_image.nb_hidden_image_index_save_or_not, ui.gallery_state.txt_gallery] - btn_load_datasets.click( - fn=load_files_from_dir, - inputs=[tb_img_directory, tb_caption_file_ext, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, dd_intterogator_names, cb_use_custom_threshold_booru, sl_custom_threshold_booru, cb_use_custom_threshold_waifu, sl_custom_threshold_waifu, cb_save_kohya_metadata, tb_metadata_output], - outputs= - [gl_dataset_images, gl_filter_images, txt_gallery, txt_selection] + - [cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply] + - o_filter_and_gallery - ) - btn_load_datasets.click( - fn=lambda:['', '', '', -1, -1, -1], - outputs=[tb_common_tags, tb_edit_tags, tb_caption_selected_image, nb_hidden_image_index, nb_hidden_image_index_prev, nb_hidden_image_index_save_or_not] - ) - - cb_save_kohya_metadata.change( - fn=lambda x:gr.update(visible=x), - inputs=cb_save_kohya_metadata, - outputs=kohya_metadata - ) - - #---------------------------------------------------------------- - # Filter by Tags tab - - tag_filter_ui.set_callbacks( - on_filter_update=lambda a, b: - update_gallery() + - update_common_tags() + - [', '.join(tag_filter_ui.filter.tags)] + - [get_current_move_or_delete_target_num(a, b)] + - [tag_select_ui_remove.cbg_tags_update()], - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs= - [cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply, nb_hidden_image_index, nb_hidden_image_index_prev, nb_hidden_image_index_save_or_not, txt_gallery] + - [tb_common_tags, tb_edit_tags] + - [tb_sr_selected_tags] + - [ta_move_or_delete_target_dataset_num]+ - [tag_select_ui_remove.cbg_tags], - _js='(...args) => {dataset_tag_editor_gl_dataset_images_close(); return args}' - ) - - tag_filter_ui_neg.set_callbacks( - on_filter_update=lambda a, b: - update_gallery() + - update_common_tags() + - [get_current_move_or_delete_target_num(a, b)] + - [tag_select_ui_remove.cbg_tags_update()], - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs= - [cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply, nb_hidden_image_index, nb_hidden_image_index_prev, nb_hidden_image_index_save_or_not, txt_gallery] + - [tb_common_tags, tb_edit_tags] + - [ta_move_or_delete_target_dataset_num] + - [tag_select_ui_remove.cbg_tags], - _js='(...args) => {dataset_tag_editor_gl_dataset_images_close(); return args}' - ) - - nb_hidden_dataset_filter_apply.change( - fn=lambda a, b: [a, b], - _js='(x, y) => [y>=0 ? dataset_tag_editor_gl_dataset_images_filter(x) : x, -1]', - inputs=[cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply], - outputs=[cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply] - ) - - btn_clear_tag_filters.click( - fn=clear_tag_filters, - outputs=o_filter_and_gallery - ) - btn_clear_tag_filters.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - - btn_clear_all_filters.click( - fn=clear_image_selection, - outputs=[gl_filter_images, txt_selection, nb_hidden_selection_image_index] - ) - btn_clear_all_filters.click( - fn=clear_tag_filters, - outputs=o_filter_and_gallery - ) - btn_clear_all_filters.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - - #---------------------------------------------------------------- - # Filter by Selection tab - - btn_hidden_set_selection_index.click( - fn=selection_index_changed, - _js="(x) => [dataset_tag_editor_gl_filter_images_selected_index()]", - inputs=[nb_hidden_selection_image_index], - outputs=[txt_selection, nb_hidden_selection_image_index] - ) - - btn_add_image_selection.click( - fn=add_image_selection, - inputs=[nb_hidden_image_index], - outputs=[gl_filter_images, nb_hidden_image_index] - ) - - btn_add_all_displayed_image_selection.click( - fn=add_all_displayed_image_selection, - outputs=gl_filter_images - ) - - btn_invert_image_selection.click( - fn=invert_image_selection, - outputs=gl_filter_images - ) - - btn_remove_image_selection.click( - fn=remove_image_selection, - inputs=[nb_hidden_selection_image_index], - outputs=[gl_filter_images, txt_selection, nb_hidden_selection_image_index] - ) - - btn_clear_image_selection.click( - fn=clear_image_selection, - outputs=[gl_filter_images, txt_selection, nb_hidden_selection_image_index] - ) - - btn_apply_image_selection_filter.click( - fn=apply_image_selection_filter, - outputs=o_filter_and_gallery - ) - btn_apply_image_selection_filter.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_apply_image_selection_filter.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - - #---------------------------------------------------------------- - # Batch Edit Captions tab - - btn_apply_edit_tags.click( - fn=apply_edit_tags, - inputs=[tb_common_tags, tb_edit_tags, cb_prepend_tags], - outputs=o_filter_and_gallery - ) - btn_apply_edit_tags.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_apply_edit_tags.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - - btn_apply_sr_tags.click( - fn=search_and_replace, - inputs=[tb_sr_search_tags, tb_sr_replace_tags, rb_sr_replace_target, cb_use_regex], - outputs=o_filter_and_gallery - ) - btn_apply_sr_tags.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_apply_sr_tags.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - - cb_show_only_tags_selected.change( - fn=cb_show_only_tags_selected_changed, - inputs=cb_show_only_tags_selected, - outputs=[tb_common_tags, tb_edit_tags] - ) - - btn_remove_duplicate.click( - fn=remove_duplicated_tags, - outputs=o_filter_and_gallery - ) - - tag_select_ui_remove.set_callbacks() - - btn_remove_selected.click( - fn=remove_selected_tags, - outputs=o_filter_and_gallery - ) - - btn_sort_selected.click( - fn=sort_selected_tags, - outputs=o_filter_and_gallery - ) - - #---------------------------------------------------------------- - # Edit Caption of Selected Image tab - - btn_hidden_set_index.click( - fn=lambda x, y:[x, y], - _js="(x, y) => [dataset_tag_editor_gl_dataset_images_selected_index(), x]", - inputs=[nb_hidden_image_index, nb_hidden_image_index_prev], - outputs=[nb_hidden_image_index, nb_hidden_image_index_prev] - ) + o_update_filter_and_gallery = \ + [ui.filter_by_tags.tag_filter_ui.cbg_tags, ui.filter_by_tags.tag_filter_ui_neg.cbg_tags] + \ + o_update_gallery + \ + [ui.batch_edit_captions.tb_common_tags, ui.batch_edit_captions.tb_edit_tags] + \ + [ui.batch_edit_captions.tb_sr_selected_tags] +\ + [ui.batch_edit_captions.tag_select_ui_remove.cbg_tags] - nb_hidden_image_index.change( - fn=gallery_index_changed, - inputs=[nb_hidden_image_index_prev, nb_hidden_image_index, tb_edit_caption_selected_image, cb_copy_caption_automatically, cb_ask_save_when_caption_changed], - outputs=[nb_hidden_image_index_save_or_not] + [tb_caption_selected_image, tb_edit_caption_selected_image] + [tb_hidden_edit_caption] + [txt_gallery] - ) - - nb_hidden_image_index_save_or_not.change( - fn=lambda a:None, - _js='(a) => dataset_tag_editor_ask_save_change_or_not(a)', - inputs=nb_hidden_image_index_save_or_not - ) - - btn_hidden_save_caption.click( - fn=lambda iprev, e, s, t, inext: dialog_selected_save_caption_change(iprev, e, s) + [get_current_move_or_delete_target_num(t, inext)], - inputs=[nb_hidden_image_index_prev, tb_hidden_edit_caption, cb_sort_caption_on_save] + [rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=o_filter_and_gallery + [ta_move_or_delete_target_dataset_num] - ) - - btn_copy_caption.click( - fn=lambda a:a, - inputs=[tb_caption_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_append_caption.click( - fn=lambda a, b : b + (', ' if a and b else '') + a, - inputs=[tb_caption_selected_image, tb_edit_caption_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_prepend_caption.click( - fn=lambda a, b : a + (', ' if a and b else '') + b, - inputs=[tb_caption_selected_image, tb_edit_caption_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_interrogate_si.click( - fn=interrogate_selected_image, - inputs=[dd_intterogator_names_si, cb_use_custom_threshold_booru, sl_custom_threshold_booru, cb_use_custom_threshold_waifu, sl_custom_threshold_waifu], - outputs=[tb_interrogate_selected_image] - ) - - btn_copy_interrogate.click( - fn=lambda a:a, - inputs=[tb_interrogate_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_append_interrogate.click( - fn=lambda a, b : b + (', ' if a and b else '') + a, - inputs=[tb_interrogate_selected_image, tb_edit_caption_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_prepend_interrogate.click( - fn=lambda a, b : a + (', ' if a and b else '') + b, - inputs=[tb_interrogate_selected_image, tb_edit_caption_selected_image], - outputs=[tb_edit_caption_selected_image] - ) - - btn_apply_changes_selected_image.click( - fn=change_selected_image_caption, - inputs=[tb_edit_caption_selected_image, nb_hidden_image_index, cb_sort_caption_on_save], - outputs=o_filter_and_gallery - ) - btn_apply_changes_selected_image.click( - fn=lambda a:a, - inputs=[tb_edit_caption_selected_image], - outputs=[tb_caption_selected_image] - ) - btn_apply_changes_selected_image.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_apply_changes_selected_image.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - - - - #---------------------------------------------------------------- - # Move or Delete Files - - rb_move_or_delete_target_data.change( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - - btn_move_or_delete_move_files.click( - fn=move_files, - inputs=[rb_move_or_delete_target_data, cbg_move_or_delete_target_file, tb_move_or_delete_caption_ext, tb_move_or_delete_destination_dir, nb_hidden_image_index], - outputs=o_filter_and_gallery - ) - btn_move_or_delete_move_files.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_move_or_delete_move_files.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - - btn_move_or_delete_delete_files.click( - fn=delete_files, - inputs=[rb_move_or_delete_target_data, cbg_move_or_delete_target_file, tb_move_or_delete_caption_ext, nb_hidden_image_index], - outputs=o_filter_and_gallery - ) - btn_move_or_delete_delete_files.click( - fn=get_current_move_or_delete_target_num, - inputs=[rb_move_or_delete_target_data, nb_hidden_image_index], - outputs=[ta_move_or_delete_target_dataset_num] - ) - btn_move_or_delete_delete_files.click( - fn=None, - _js='() => dataset_tag_editor_gl_dataset_images_close()' - ) - + ui.toprow.set_callbacks(ui.load_dataset) + ui.load_dataset.set_callbacks(o_update_filter_and_gallery,ui.toprow, ui.dataset_gallery, ui.filter_by_tags, ui.filter_by_selection, ui.batch_edit_captions, update_filter_and_gallery) + ui.dataset_gallery.set_callbacks(ui.load_dataset, ui.gallery_state, get_filters) + ui.gallery_state.set_callbacks(ui.dataset_gallery) + ui.filter_by_tags.set_callbacks(o_update_gallery, o_update_filter_and_gallery, ui.batch_edit_captions, ui.move_or_delete_files, update_gallery, update_filter_and_gallery, get_filters) + ui.filter_by_selection.set_callbacks(o_update_filter_and_gallery, ui.dataset_gallery, ui.filter_by_tags, get_filters, update_filter_and_gallery) + ui.batch_edit_captions.set_callbacks(o_update_filter_and_gallery, ui.load_dataset, ui.filter_by_tags, get_filters, update_filter_and_gallery) + ui.edit_caption_of_selected_image.set_callbacks(o_update_filter_and_gallery, ui.dataset_gallery, ui.load_dataset, get_filters, update_filter_and_gallery) + ui.move_or_delete_files.set_callbacks(o_update_filter_and_gallery, ui.dataset_gallery, ui.filter_by_tags, ui.batch_edit_captions, ui.filter_by_selection, ui.edit_caption_of_selected_image, get_filters, update_filter_and_gallery) + return [(dataset_tag_editor_interface, "Dataset Tag Editor", "dataset_tag_editor_interface")] diff --git a/scripts/ui/__init__.py b/scripts/ui/__init__.py new file mode 100644 index 0000000..07fff21 --- /dev/null +++ b/scripts/ui/__init__.py @@ -0,0 +1,2 @@ +from .ui_common import * +from .ui_instance import * \ No newline at end of file diff --git a/scripts/ui/block_dataset_gallery.py b/scripts/ui/block_dataset_gallery.py new file mode 100644 index 0000000..73daeab --- /dev/null +++ b/scripts/ui/block_dataset_gallery.py @@ -0,0 +1,60 @@ +import gradio as gr + +from .ui_common import * +from .uibase import UIBase + +class DatasetGalleryUI(UIBase): + def __init__(self): + self.selected_path = '' + self.selected_index = -1 + self.selected_index_prev = -1 + + def create_ui(self, image_columns): + with gr.Row(visible=False): + self.btn_hidden_set_index = gr.Button(elem_id="dataset_tag_editor_btn_hidden_set_index") + self.nb_hidden_image_index = gr.Number(value=None, label='hidden_idx_next') + self.nb_hidden_image_index_prev = gr.Number(value=None, label='hidden_idx_prev') + self.gl_dataset_images = gr.Gallery(label='Dataset Images', elem_id="dataset_tag_editor_dataset_gallery").style(grid=image_columns) + + def set_callbacks(self, load_dataset, gallery_state, get_filters): + gallery_state.register_value('Selected Image', self.selected_path) + + load_dataset.btn_load_datasets.click( + fn=lambda:[-1, -1], + outputs=[self.nb_hidden_image_index, self.nb_hidden_image_index_prev] + ) + + def set_index_clicked(next_idx: int, prev_idx: int): + prev_idx = int(prev_idx) if prev_idx is not None else -1 + next_idx = int(next_idx) if next_idx is not None else -1 + img_paths = dte_instance.get_filtered_imgpaths(filters=get_filters()) + + if prev_idx < 0 or len(img_paths) <= prev_idx: + prev_idx = -1 + + if 0 <= next_idx and next_idx < len(img_paths): + self.selected_path = img_paths[next_idx] + else: + next_idx = -1 + self.selected_path = '' + + gallery_state.register_value('Selected Image', self.selected_path) + return [next_idx, prev_idx] + + self.btn_hidden_set_index.click( + fn=set_index_clicked, + _js="(x, y) => [dataset_tag_editor_gl_dataset_images_selected_index(), x]", + inputs=[self.nb_hidden_image_index, self.nb_hidden_image_index_prev], + outputs=[self.nb_hidden_image_index, self.nb_hidden_image_index_prev] + ) + self.nb_hidden_image_index.change( + fn=self.func_to_set_value('selected_index', int), + inputs=self.nb_hidden_image_index + ) + self.nb_hidden_image_index_prev.change( + fn=self.func_to_set_value('selected_index_prev', int), + inputs=self.nb_hidden_image_index_prev + ) + + + diff --git a/scripts/ui/block_gallery_state.py b/scripts/ui/block_gallery_state.py new file mode 100644 index 0000000..cab3013 --- /dev/null +++ b/scripts/ui/block_gallery_state.py @@ -0,0 +1,35 @@ +import gradio as gr + +from .ui_common import * +from .uibase import UIBase + +class GalleryStateUI(UIBase): + def __init__(self): + self.texts = dict() + + def register_value(self, key:str, value:str): + self.texts[key] = value + + def remove_key(self, key:str): + del self.texts[key] + + def get_current_gallery_txt(self): + res = '' + for k, v in self.texts.items(): + res += f'{k} : {v}
' + return res + + def create_ui(self): + self.txt_gallery = gr.HTML(value=self.get_current_gallery_txt()) + + def set_callbacks(self, dataset_gallery): + dataset_gallery.nb_hidden_image_index.change( + fn=self.update_gallery_txt, + inputs=None, + outputs=self.txt_gallery + ) + + def update_gallery_txt(self): + return self.get_current_gallery_txt() + + diff --git a/scripts/ui/block_load_dataset.py b/scripts/ui/block_load_dataset.py new file mode 100644 index 0000000..60baba8 --- /dev/null +++ b/scripts/ui/block_load_dataset.py @@ -0,0 +1,95 @@ +from typing import List +import gradio as gr + +from modules import shared +from modules.shared import opts + +from .ui_common import * +from .uibase import UIBase + + +INTERROGATOR_NAMES = dte_module.INTERROGATOR_NAMES +InterrogateMethod = dte_module.InterrogateMethod + + +class LoadDatasetUI(UIBase): + def __init__(self): + self.caption_file_ext = '' + + def create_ui(self, cfg_general): + with gr.Column(variant='panel'): + with gr.Row(): + with gr.Column(scale=3): + self.tb_img_directory = gr.Textbox(label='Dataset directory', placeholder='C:\\directory\\of\\datasets', value=cfg_general.dataset_dir) + with gr.Column(scale=1, min_width=60): + self.tb_caption_file_ext = gr.Textbox(label='Caption File Ext', placeholder='.txt (on Load and Save)', value=cfg_general.caption_ext) + self.caption_file_ext = cfg_general.caption_ext + with gr.Column(scale=1, min_width=80): + self.btn_load_datasets = gr.Button(value='Load') + with gr.Accordion(label='Dataset Load Settings'): + with gr.Row(): + with gr.Column(): + self.cb_load_recursive = gr.Checkbox(value=cfg_general.load_recursive, label='Load from subdirectories') + self.cb_load_caption_from_filename = gr.Checkbox(value=cfg_general.load_caption_from_filename, label='Load caption from filename if no text file exists') + with gr.Column(): + self.rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite', 'Prepend', 'Append'], value=cfg_general.use_interrogator, label='Use Interrogator Caption') + self.dd_intterogator_names = gr.Dropdown(label = 'Interrogators', choices=INTERROGATOR_NAMES, value=cfg_general.use_interrogator_names, interactive=True, multiselect=True) + with gr.Accordion(label='Interrogator Settings', open=False): + with gr.Row(): + self.cb_use_custom_threshold_booru = gr.Checkbox(value=cfg_general.use_custom_threshold_booru, label='Use Custom Threshold (Booru)', interactive=True) + self.sl_custom_threshold_booru = gr.Slider(minimum=0, maximum=1, value=cfg_general.custom_threshold_booru, step=0.01, interactive=True, label='Booru Score Threshold') + with gr.Row(): + self.cb_use_custom_threshold_waifu = gr.Checkbox(value=cfg_general.use_custom_threshold_waifu, label='Use Custom Threshold (WDv1.4 Tagger)', interactive=True) + self.sl_custom_threshold_waifu = gr.Slider(minimum=0, maximum=1, value=cfg_general.custom_threshold_waifu, step=0.01, interactive=True, label='WDv1.4 Tagger Score Threshold') + + def set_callbacks(self, o_update_filter_and_gallery, toprow, dataset_gallery, filter_by_tags, filter_by_selection, batch_edit_captions, update_filter_and_gallery): + def load_files_from_dir( + dir: str, + caption_file_ext: str, + recursive: bool, + load_caption_from_filename: bool, + use_interrogator: str, + use_interrogator_names: List[str], + use_custom_threshold_booru: bool, + custom_threshold_booru: float, + use_custom_threshold_waifu: bool, + custom_threshold_waifu: float, + use_kohya_metadata: bool, + kohya_json_path: str + ): + + interrogate_method = InterrogateMethod.NONE + if use_interrogator == 'If Empty': + 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 + + threshold_booru = custom_threshold_booru if use_custom_threshold_booru else shared.opts.interrogate_deepbooru_score_threshold + threshold_waifu = custom_threshold_waifu if use_custom_threshold_waifu else -1 + + dte_instance.load_dataset(dir, caption_file_ext, recursive, load_caption_from_filename, interrogate_method, use_interrogator_names, threshold_booru, threshold_waifu, opts.dataset_editor_use_temp_files, kohya_json_path if use_kohya_metadata else None) + imgs = dte_instance.get_filtered_imgs(filters=[]) + img_indices = dte_instance.get_filtered_imgindices(filters=[]) + return [ + imgs, + [] + ] +\ + [gr.CheckboxGroup.update(value=[str(i) for i in img_indices], choices=[str(i) for i in img_indices]), True] +\ + filter_by_tags.clear_filters(update_filter_and_gallery) +\ + [batch_edit_captions.tag_select_ui_remove.cbg_tags_update()] + + self.btn_load_datasets.click( + fn=load_files_from_dir, + inputs=[self.tb_img_directory, self.tb_caption_file_ext, self.cb_load_recursive, self.cb_load_caption_from_filename, self.rb_use_interrogator, self.dd_intterogator_names, self.cb_use_custom_threshold_booru, self.sl_custom_threshold_booru, self.cb_use_custom_threshold_waifu, self.sl_custom_threshold_waifu, toprow.cb_save_kohya_metadata, toprow.tb_metadata_output], + outputs= + [dataset_gallery.gl_dataset_images, filter_by_selection.gl_filter_images] + + [filter_by_tags.cbg_hidden_dataset_filter, filter_by_tags.nb_hidden_dataset_filter_apply] + + o_update_filter_and_gallery + ) + + + diff --git a/scripts/ui/block_tag_filter.py b/scripts/ui/block_tag_filter.py new file mode 100644 index 0000000..91aca91 --- /dev/null +++ b/scripts/ui/block_tag_filter.py @@ -0,0 +1,134 @@ +from typing import List, Callable +import gradio as gr + +from .ui_common import * + +filters = dte_module.filters +TagFilter = filters.TagFilter + +class TagFilterUI(): + def __init__(self, tag_filter_mode = TagFilter.Mode.INCLUSIVE): + self.logic = TagFilter.Logic.AND + self.filter_word = '' + self.sort_by = 'Alphabetical Order' + self.sort_order = 'Ascending' + self.selected_tags = set() + self.filter_mode = tag_filter_mode + self.filter = TagFilter(logic=self.logic, mode=self.filter_mode) + self.get_filters = lambda:[] + self.prefix = False + self.suffix = False + self.regex = False + self.on_filter_update_callbacks = [] + + def get_filter(self): + return self.filter + + def create_ui(self, get_filters: Callable[[], List[filters.Filter]], logic = TagFilter.Logic.AND, sort_by = 'Alphabetical Order', sort_order = 'Ascending', prefix=False, suffix=False, regex=False): + self.get_filters = get_filters + self.logic = logic + self.filter = filters.TagFilter(logic=self.logic, mode=self.filter_mode) + self.sort_by = sort_by + self.sort_order = sort_order + self.prefix = prefix + self.suffix = suffix + self.regex = regex + + self.tb_search_tags = gr.Textbox(label='Search Tags', interactive=True) + with gr.Row(): + self.cb_prefix = gr.Checkbox(label='Prefix', value=self.prefix, interactive=True) + self.cb_suffix = gr.Checkbox(label='Suffix', value=self.suffix, interactive=True) + self.cb_regex = gr.Checkbox(label='Use regex', value=self.regex, interactive=True) + with gr.Row(): + self.rb_sort_by = gr.Radio(choices=['Alphabetical Order', 'Frequency', 'Length'], value=sort_by, interactive=True, label='Sort by') + self.rb_sort_order = gr.Radio(choices=['Ascending', 'Descending'], value=sort_order, interactive=True, label='Sort Order') + v = 'AND' if self.logic==TagFilter.Logic.AND else 'OR' if self.logic==TagFilter.Logic.OR else 'NONE' + self.rb_logic = gr.Radio(choices=['AND', 'OR', 'NONE'], value=v, label='Filter Logic', interactive=True) + self.cbg_tags = gr.CheckboxGroup(label='Filter Images by Tags', interactive=True) + + + def on_filter_update(self, fn:Callable[[List], List], inputs=None, outputs=None, _js=None): + self.on_filter_update_callbacks.append((fn, inputs, outputs, _js)) + + + def set_callbacks(self): + self.tb_search_tags.change(fn=self.tb_search_tags_changed, inputs=self.tb_search_tags, outputs=self.cbg_tags) + self.cb_prefix.change(fn=self.cb_prefix_changed, inputs=self.cb_prefix, outputs=self.cbg_tags) + self.cb_suffix.change(fn=self.cb_suffix_changed, inputs=self.cb_suffix, outputs=self.cbg_tags) + self.cb_regex.change(fn=self.cb_regex_changed, inputs=self.cb_regex, outputs=self.cbg_tags) + self.rb_sort_by.change(fn=self.rd_sort_by_changed, inputs=self.rb_sort_by, outputs=self.cbg_tags) + self.rb_sort_order.change(fn=self.rd_sort_order_changed, inputs=self.rb_sort_order, outputs=self.cbg_tags) + + self.rb_logic.change(fn=self.rd_logic_changed, inputs=[self.rb_logic], outputs=[self.cbg_tags]) + for fn, inputs, outputs, _js in self.on_filter_update_callbacks: + self.rb_logic.change(fn=fn, inputs=inputs, outputs=outputs, _js=_js) + self.cbg_tags.change(fn=self.cbg_tags_changed, inputs=[self.cbg_tags], outputs=[self.cbg_tags]) + for fn, inputs, outputs, _js in self.on_filter_update_callbacks: + self.cbg_tags.change(fn=fn, inputs=inputs, outputs=outputs, _js=_js) + + + def tb_search_tags_changed(self, tb_search_tags: str): + self.filter_word = tb_search_tags + return self.cbg_tags_update() + + + def cb_prefix_changed(self, prefix:bool): + self.prefix = prefix + return self.cbg_tags_update() + + + def cb_suffix_changed(self, suffix:bool): + self.suffix = suffix + return self.cbg_tags_update() + + + def cb_regex_changed(self, use_regex:bool): + self.regex = use_regex + return self.cbg_tags_update() + + + def rd_sort_by_changed(self, rb_sort_by: str): + self.sort_by = rb_sort_by + return self.cbg_tags_update() + + + def rd_sort_order_changed(self, rd_sort_order: str): + self.sort_order = rd_sort_order + return self.cbg_tags_update() + + + def rd_logic_changed(self, rd_logic: str): + self.logic = TagFilter.Logic.AND if rd_logic == 'AND' else TagFilter.Logic.OR if rd_logic == 'OR' else TagFilter.Logic.NONE + self.filter = TagFilter(self.selected_tags, self.logic, self.filter_mode) + return self.cbg_tags_update() + + + def cbg_tags_changed(self, cbg_tags: List[str]): + self.selected_tags = dte_instance.cleanup_tagset(set(dte_instance.read_tags(cbg_tags))) + return self.cbg_tags_update() + + + def cbg_tags_update(self): + self.selected_tags = dte_instance.cleanup_tagset(self.selected_tags) + self.filter = TagFilter(self.selected_tags, self.logic, self.filter_mode) + + if self.filter_mode == TagFilter.Mode.INCLUSIVE: + tags = dte_instance.get_filtered_tags(self.get_filters(), self.filter_word, self.filter.logic == TagFilter.Logic.AND, prefix=self.prefix, suffix=self.suffix, regex=self.regex) + else: + tags = dte_instance.get_filtered_tags(self.get_filters(), self.filter_word, self.filter.logic == TagFilter.Logic.OR, prefix=self.prefix, suffix=self.suffix, regex=self.regex) + tags_in_filter = self.filter.tags + + tags = dte_instance.sort_tags(tags=tags, sort_by=self.sort_by, sort_order=self.sort_order) + tags_in_filter = dte_instance.sort_tags(tags=tags_in_filter, sort_by=self.sort_by, sort_order=self.sort_order) + + tags = tags_in_filter + [tag for tag in tags if tag not in self.filter.tags] + tags = dte_instance.write_tags(tags) + tags_in_filter = dte_instance.write_tags(tags_in_filter) + + return gr.CheckboxGroup.update(value=tags_in_filter, choices=tags) + + + def clear_filter(self): + self.filter = TagFilter(logic=self.logic, mode=self.filter_mode) + self.filter_word = '' + self.selected_tags = set() \ No newline at end of file diff --git a/scripts/ui/block_tag_select.py b/scripts/ui/block_tag_select.py new file mode 100644 index 0000000..47e4b42 --- /dev/null +++ b/scripts/ui/block_tag_select.py @@ -0,0 +1,110 @@ +from typing import List, Callable +import gradio as gr + +from .ui_common import * + +TagFilter = dte_module.filters.TagFilter +Filter = dte_module.filters.Filter + + +class TagSelectUI(): + def __init__(self): + self.filter_word = '' + self.sort_by = 'Alphabetical Order' + self.sort_order = 'Ascending' + self.selected_tags = set() + self.tags = set() + self.get_filters = lambda:[] + self.prefix = False + self.suffix = False + self.regex = False + + + def create_ui(self, get_filters: Callable[[], List[Filter]], sort_by = 'Alphabetical Order', sort_order = 'Ascending', prefix=False, suffix=False, regex=False): + self.get_filters = get_filters + self.prefix = prefix + self.suffix = suffix + self.regex = regex + + self.tb_search_tags = gr.Textbox(label='Search Tags', interactive=True) + with gr.Row(): + self.cb_prefix = gr.Checkbox(label='Prefix', value=False, interactive=True) + self.cb_suffix = gr.Checkbox(label='Suffix', value=False, interactive=True) + self.cb_regex = gr.Checkbox(label='Use regex', value=False, interactive=True) + with gr.Row(): + self.rb_sort_by = gr.Radio(choices=['Alphabetical Order', 'Frequency', 'Length'], value=sort_by, interactive=True, label='Sort by') + self.rb_sort_order = gr.Radio(choices=['Ascending', 'Descending'], value=sort_order, interactive=True, label='Sort Order') + with gr.Row(): + self.btn_select_visibles = gr.Button(value='Select visible tags') + self.btn_deselect_visibles = gr.Button(value='Deselect visible tags') + self.cbg_tags = gr.CheckboxGroup(label='Select Tags', interactive=True) + + + def set_callbacks(self): + self.tb_search_tags.change(fn=self.tb_search_tags_changed, inputs=self.tb_search_tags, outputs=self.cbg_tags) + self.cb_prefix.change(fn=self.cb_prefix_changed, inputs=self.cb_prefix, outputs=self.cbg_tags) + self.cb_suffix.change(fn=self.cb_suffix_changed, inputs=self.cb_suffix, outputs=self.cbg_tags) + self.cb_regex.change(fn=self.cb_regex_changed, inputs=self.cb_regex, outputs=self.cbg_tags) + self.rb_sort_by.change(fn=self.rd_sort_by_changed, inputs=self.rb_sort_by, outputs=self.cbg_tags) + self.rb_sort_order.change(fn=self.rd_sort_order_changed, inputs=self.rb_sort_order, outputs=self.cbg_tags) + self.btn_select_visibles.click(fn=self.btn_select_visibles_clicked, outputs=self.cbg_tags) + self.btn_deselect_visibles.click(fn=self.btn_deselect_visibles_clicked, inputs=self.cbg_tags, outputs=self.cbg_tags) + self.cbg_tags.change(fn=self.cbg_tags_changed, inputs=self.cbg_tags, outputs=self.cbg_tags) + + + def tb_search_tags_changed(self, tb_search_tags: str): + self.filter_word = tb_search_tags + return self.cbg_tags_update() + + + def cb_prefix_changed(self, prefix:bool): + self.prefix = prefix + return self.cbg_tags_update() + + + def cb_suffix_changed(self, suffix:bool): + self.suffix = suffix + return self.cbg_tags_update() + + + def cb_regex_changed(self, regex:bool): + self.regex = regex + return self.cbg_tags_update() + + + def rd_sort_by_changed(self, rb_sort_by: str): + self.sort_by = rb_sort_by + return self.cbg_tags_update() + + + def rd_sort_order_changed(self, rd_sort_order: str): + self.sort_order = rd_sort_order + return self.cbg_tags_update() + + + def cbg_tags_changed(self, cbg_tags: List[str]): + self.selected_tags = set(dte_instance.read_tags(cbg_tags)) + return self.cbg_tags_update() + + + def btn_deselect_visibles_clicked(self, cbg_tags: List[str]): + tags = dte_instance.get_filtered_tags(self.get_filters(), self.filter_word, True) + selected_tags = set(dte_instance.read_tags(cbg_tags)) & tags + self.selected_tags -= selected_tags + return self.cbg_tags_update() + + + def btn_select_visibles_clicked(self): + tags = set(dte_instance.get_filtered_tags(self.get_filters(), self.filter_word, True)) + self.selected_tags |= tags + return self.cbg_tags_update() + + + def cbg_tags_update(self): + tags = dte_instance.get_filtered_tags(self.get_filters(), self.filter_word, True, prefix=self.prefix, suffix=self.suffix, regex=self.regex) + self.tags = set(dte_instance.get_filtered_tags(self.get_filters(), filter_tags=True, prefix=self.prefix, suffix=self.suffix, regex=self.regex)) + self.selected_tags &= self.tags + tags = dte_instance.sort_tags(tags=tags, sort_by=self.sort_by, sort_order=self.sort_order) + tags = dte_instance.write_tags(tags) + selected_tags = dte_instance.write_tags(list(self.selected_tags)) + return gr.CheckboxGroup.update(value=selected_tags, choices=tags) \ No newline at end of file diff --git a/scripts/ui/block_toprow.py b/scripts/ui/block_toprow.py new file mode 100644 index 0000000..738e07c --- /dev/null +++ b/scripts/ui/block_toprow.py @@ -0,0 +1,46 @@ +import gradio as gr + +from .ui_common import * +from .uibase import UIBase + + +class ToprowUI(UIBase): + def create_ui(self, cfg_general): + with gr.Column(variant='panel'): + with gr.Row(): + with gr.Column(scale=1): + self.btn_save_all_changes = gr.Button(value='Save all changes', variant='primary') + with gr.Column(scale=2): + self.cb_backup = gr.Checkbox(value=cfg_general.backup, label='Backup original text file (original file will be renamed like filename.000, .001, .002, ...)', interactive=True) + gr.HTML(value='Note: New text file will be created if you are using filename as captions.') + with gr.Row(): + self.cb_save_kohya_metadata = gr.Checkbox(value=cfg_general.save_kohya_metadata, label="Use kohya-ss's finetuning metadata json", interactive=True) + with gr.Row(): + with gr.Column(variant='panel', visible=cfg_general.save_kohya_metadata) as self.cl_kohya_metadata: + self.tb_metadata_output = gr.Textbox(label='json path', placeholder='C:\\path\\to\\metadata.json',value=cfg_general.meta_output_path) + self.tb_metadata_input = gr.Textbox(label='json input path (Optional, only for append results)', placeholder='C:\\path\\to\\metadata.json',value=cfg_general.meta_input_path) + with gr.Row(): + self.cb_metadata_overwrite = gr.Checkbox(value=cfg_general.meta_overwrite, label="Overwrite if output file exists", interactive=True) + self.cb_metadata_as_caption = gr.Checkbox(value=cfg_general.meta_save_as_caption, label="Save metadata as caption", interactive=True) + self.cb_metadata_use_fullpath = gr.Checkbox(value=cfg_general.meta_use_full_path, label="Save metadata image key as fullpath", interactive=True) + with gr.Row(visible=False): + self.txt_result = gr.Textbox(label='Results', interactive=False) + + def set_callbacks(self, load_dataset): + + def save_all_changes(backup: bool, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool): + if not metadata_input: + metadata_input = None + dte_instance.save_dataset(backup, load_dataset.caption_file_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath) + + self.btn_save_all_changes.click( + fn=save_all_changes, + inputs=[self.cb_backup, self.cb_save_kohya_metadata, self.tb_metadata_output, self.tb_metadata_input, self.cb_metadata_overwrite, self.cb_metadata_as_caption, self.cb_metadata_use_fullpath] + ) + + self.cb_save_kohya_metadata.change( + fn=lambda x:gr.update(visible=x), + inputs=self.cb_save_kohya_metadata, + outputs=self.cl_kohya_metadata + ) + diff --git a/scripts/ui/tab_batch_edit_captions.py b/scripts/ui/tab_batch_edit_captions.py new file mode 100644 index 0000000..274036a --- /dev/null +++ b/scripts/ui/tab_batch_edit_captions.py @@ -0,0 +1,168 @@ +import gradio as gr + +from .ui_common import * +from .uibase import UIBase +from .block_tag_select import TagSelectUI + +class BatchEditCaptionsUI(UIBase): + def __init__(self): + self.tag_select_ui_remove = TagSelectUI() + self.show_only_selected_tags = False + + def create_ui(self, cfg_batch_edit, get_filters): + with gr.Tab(label='Search and Replace'): + with gr.Column(variant='panel'): + gr.HTML('Edit common tags.') + self.cb_show_only_tags_selected = gr.Checkbox(value=cfg_batch_edit.show_only_selected, label='Show only the tags selected in the Positive Filter') + self.show_only_selected_tags = cfg_batch_edit.show_only_selected + self.tb_common_tags = gr.Textbox(label='Common Tags', interactive=False) + self.tb_edit_tags = gr.Textbox(label='Edit Tags', interactive=True) + self.cb_prepend_tags = gr.Checkbox(value=cfg_batch_edit.prepend, label='Prepend additional tags') + self.btn_apply_edit_tags = gr.Button(value='Apply changes to filtered images', variant='primary') + with gr.Accordion(label='Show description of how to edit tags', open=False): + gr.HTML(value=""" + 1. The tags common to all displayed images are shown in comma separated style.
+ 2. When changes are applied, all tags in each displayed images are replaced.
+ 3. If you change some tags into blank, they will be erased.
+ 4. If you add some tags to the end, they will be added to the end/beginning of the text file.
+ 5. Changes are not applied to the text files until the "Save all changes" button is pressed.
+ ex A.
+  Original Text = "A, A, B, C" Common Tags = "B, A" Edit Tags = "X, Y"
+  Result = "Y, Y, X, C" (B->X, A->Y)
+ ex B.
+  Original Text = "A, B, C" Common Tags = "(nothing)" Edit Tags = "X, Y"
+  Result = "A, B, C, X, Y" (add X and Y to the end (default))
+  Result = "X, Y, A, B, C" (add X and Y to the beginning ("Prepend additional tags" checked))
+ ex C.
+  Original Text = "A, B, C, D, E" Common Tags = "A, B, D" Edit Tags = ", X, "
+  Result = "X, C, E" (A->"", B->X, D->"")
+ """) + with gr.Column(variant='panel'): + gr.HTML('Search and Replace for all images displayed.') + self.tb_sr_search_tags = gr.Textbox(label='Search Text', interactive=True) + self.tb_sr_replace_tags = gr.Textbox(label='Replace Text', interactive=True) + self.cb_use_regex = gr.Checkbox(label='Use regex', value=cfg_batch_edit.use_regex) + self.rb_sr_replace_target = gr.Radio(['Only Selected Tags', 'Each Tags', 'Entire Caption'], value=cfg_batch_edit.target, label='Search and Replace in', interactive=True) + self.tb_sr_selected_tags = gr.Textbox(label='Selected Tags', interactive=False, lines=2) + self.btn_apply_sr_tags = gr.Button(value='Search and Replace', variant='primary') + with gr.Tab(label='Remove'): + with gr.Column(variant='panel'): + gr.HTML('Remove duplicate tags from the images displayed.') + self.btn_remove_duplicate = gr.Button(value='Remove duplicate tags', variant='primary') + with gr.Column(variant='panel'): + gr.HTML('Remove selected tags from the images displayed.') + self.btn_remove_selected = gr.Button(value='Remove selected tags', variant='primary') + self.tag_select_ui_remove.create_ui(get_filters, cfg_batch_edit.sory_by, cfg_batch_edit.sort_order, cfg_batch_edit.sw_prefix, cfg_batch_edit.sw_suffix, cfg_batch_edit.sw_regex) + with gr.Tab(label='Extras'): + with gr.Column(variant='panel'): + gr.HTML('Sort tags from the images displayed.') + self.btn_sort_selected = gr.Button(value='Sort selected tags', variant='primary') + + def set_callbacks(self, o_update_filter_and_gallery, load_dataset, filter_by_tags, get_filters, update_filter_and_gallery): + load_dataset.btn_load_datasets.click( + fn=lambda:['', ''], + outputs=[self.tb_common_tags, self.tb_edit_tags] + ) + + def apply_edit_tags(search_tags: str, replace_tags: str, prepend: bool): + search_tags = [t.strip() for t in search_tags.split(',')] + search_tags = [t for t in search_tags if t] + replace_tags = [t.strip() for t in replace_tags.split(',')] + replace_tags = [t for t in replace_tags if t] + + dte_instance.replace_tags(search_tags = search_tags, replace_tags = replace_tags, filters=get_filters(), prepend = prepend) + filter_by_tags.tag_filter_ui.get_filter().tags = dte_instance.get_replaced_tagset(filter_by_tags.tag_filter_ui.get_filter().tags, search_tags, replace_tags) + filter_by_tags.tag_filter_ui_neg.get_filter().tags = dte_instance.get_replaced_tagset(filter_by_tags.tag_filter_ui_neg.get_filter().tags, search_tags, replace_tags) + + return update_filter_and_gallery() + + self.btn_apply_edit_tags.click( + fn=apply_edit_tags, + inputs=[self.tb_common_tags, self.tb_edit_tags, self.cb_prepend_tags], + outputs=o_update_filter_and_gallery + ) + self.btn_apply_edit_tags.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + + def search_and_replace(search_text: str, replace_text: str, target_text: str, use_regex: bool): + if target_text == 'Only Selected Tags': + selected_tags = set(filter_by_tags.tag_filter_ui.selected_tags) + dte_instance.search_and_replace_selected_tags(search_text = search_text, replace_text=replace_text, selected_tags=selected_tags, filters=get_filters(), use_regex=use_regex) + filter_by_tags.tag_filter_ui.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui.filter.tags, selected_tags, use_regex) + filter_by_tags.tag_filter_ui_neg.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui_neg.filter.tags, selected_tags, use_regex) + + elif target_text == 'Each Tags': + dte_instance.search_and_replace_selected_tags(search_text = search_text, replace_text=replace_text, selected_tags=None, filters=get_filters(), use_regex=use_regex) + filter_by_tags.tag_filter_ui.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui.filter.tags, None, use_regex) + filter_by_tags.tag_filter_ui_neg.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui_neg.filter.tags, None, use_regex) + + elif target_text == 'Entire Caption': + dte_instance.search_and_replace_caption(search_text=search_text, replace_text=replace_text, filters=get_filters(), use_regex=use_regex) + filter_by_tags.tag_filter_ui.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui.filter.tags, None, use_regex) + filter_by_tags.tag_filter_ui_neg.filter.tags = dte_instance.search_and_replace_tag_set(search_text, replace_text, filter_by_tags.tag_filter_ui_neg.filter.tags, None, use_regex) + + return update_filter_and_gallery() + + self.btn_apply_sr_tags.click( + fn=search_and_replace, + inputs=[self.tb_sr_search_tags, self.tb_sr_replace_tags, self.rb_sr_replace_target, self.cb_use_regex], + outputs=o_update_filter_and_gallery + ) + self.btn_apply_sr_tags.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + + def cb_show_only_tags_selected_changed(value: bool): + self.show_only_selected_tags = value + return self.get_common_tags(get_filters, filter_by_tags) + + self.cb_show_only_tags_selected.change( + fn=cb_show_only_tags_selected_changed, + inputs=self.cb_show_only_tags_selected, + outputs=[self.tb_common_tags, self.tb_edit_tags] + ) + + def remove_duplicated_tags(): + dte_instance.remove_duplicated_tags(get_filters()) + return update_filter_and_gallery() + + self.btn_remove_duplicate.click( + fn=remove_duplicated_tags, + outputs=o_update_filter_and_gallery + ) + + self.tag_select_ui_remove.set_callbacks() + + def remove_selected_tags(): + dte_instance.remove_tags(self.tag_select_ui_remove.selected_tags, get_filters()) + return update_filter_and_gallery() + + self.btn_remove_selected.click( + fn=remove_selected_tags, + outputs=o_update_filter_and_gallery + ) + + def sort_selected_tags(**sort_args): + dte_instance.sort_filtered_tags(get_filters(), **sort_args) + return update_filter_and_gallery() + + self.btn_sort_selected.click( + fn=sort_selected_tags, + outputs=o_update_filter_and_gallery + ) + + self.cb_show_only_tags_selected.change( + fn=self.func_to_set_value('show_only_selected_tags'), + inputs=self.cb_show_only_tags_selected + ) + + + def get_common_tags(self, get_filters, filter_by_tags): + if self.show_only_selected_tags: + tags = ', '.join([t for t in dte_instance.get_common_tags(filters=get_filters()) if t in filter_by_tags.tag_filter_ui.filter.tags]) + else: + tags = ', '.join(dte_instance.get_common_tags(filters=get_filters())) + return [tags, tags] \ No newline at end of file diff --git a/scripts/ui/tab_edit_caption_of_selected_image.py b/scripts/ui/tab_edit_caption_of_selected_image.py new file mode 100644 index 0000000..64b36d4 --- /dev/null +++ b/scripts/ui/tab_edit_caption_of_selected_image.py @@ -0,0 +1,162 @@ +import gradio as gr + +from modules import shared +from scripts.dte_instance import dte_module + +from .ui_common import * +from .uibase import UIBase + +class EditCaptionOfSelectedImageUI(UIBase): + def create_ui(self, cfg_edit_selected): + with gr.Row(visible=False): + self.nb_hidden_image_index_save_or_not = gr.Number(value=-1, label='hidden_s_or_n') + self.tb_hidden_edit_caption = gr.Textbox() + self.btn_hidden_save_caption = gr.Button(elem_id="dataset_tag_editor_btn_hidden_save_caption") + with gr.Tab(label='Read Caption from Selected Image'): + self.tb_caption = gr.Textbox(label='Caption of Selected Image', interactive=True, lines=6) + with gr.Row(): + self.btn_copy_caption = gr.Button(value='Copy and Overwrite') + self.btn_prepend_caption = gr.Button(value='Prepend') + self.btn_append_caption = gr.Button(value='Append') + + with gr.Tab(label='Interrogate Selected Image'): + with gr.Row(): + self.dd_intterogator_names_si = gr.Dropdown(label = 'Interrogator', choices=dte_module.INTERROGATOR_NAMES, value=cfg_edit_selected.use_interrogator_name, interactive=True, multiselect=False) + self.btn_interrogate_si = gr.Button(value='Interrogate') + self.tb_interrogate_selected_image = gr.Textbox(label='Interrogate Result', interactive=True, lines=6) + with gr.Row(): + self.btn_copy_interrogate = gr.Button(value='Copy and Overwrite') + self.btn_prepend_interrogate = gr.Button(value='Prepend') + self.btn_append_interrogate = gr.Button(value='Append') + self.cb_copy_caption_automatically = gr.Checkbox(value=cfg_edit_selected.auto_copy, label='Copy caption from selected images automatically') + self.cb_sort_caption_on_save = gr.Checkbox(value=cfg_edit_selected.sort_on_save, label='Sort caption on save') + self.cb_ask_save_when_caption_changed = gr.Checkbox(value=cfg_edit_selected.warn_change_not_saved, label='Warn if changes in caption is not saved') + self.tb_edit_caption = gr.Textbox(label='Edit Caption', interactive=True, lines=6) + self.btn_apply_changes_selected_image = gr.Button(value='Apply changes to selected image', variant='primary') + + gr.HTML("""Changes are not applied to the text files until the "Save all changes" button is pressed.""") + + def set_callbacks(self, o_update_filter_and_gallery, dataset_gallery, load_dataset, get_filters, update_filter_and_gallery): + load_dataset.btn_load_datasets.click( + fn=lambda:['', -1], + outputs=[self.tb_caption, self.nb_hidden_image_index_save_or_not] + ) + + def gallery_index_changed(next_idx:int, edit_caption: str, copy_automatically: bool, warn_change_not_saved: bool): + prev_idx = dataset_gallery.selected_index + next_idx = int(next_idx) + img_paths = dte_instance.get_filtered_imgpaths(filters=get_filters()) + prev_tags_txt = '' + if 0 <= prev_idx and prev_idx < len(img_paths): + prev_tags_txt = ', '.join(dte_instance.get_tags_by_image_path(img_paths[prev_idx])) + else: + prev_idx = -1 + + next_tags_txt = '' + if 0 <= next_idx and next_idx < len(img_paths): + next_tags_txt = ', '.join(dte_instance.get_tags_by_image_path(img_paths[next_idx])) + + return\ + [prev_idx if warn_change_not_saved and edit_caption != prev_tags_txt else -1] +\ + [next_tags_txt, next_tags_txt if copy_automatically else edit_caption] +\ + [edit_caption] + + self.nb_hidden_image_index_save_or_not.change( + fn=lambda a:None, + _js='(a) => dataset_tag_editor_ask_save_change_or_not(a)', + inputs=self.nb_hidden_image_index_save_or_not + ) + dataset_gallery.btn_hidden_set_index.click( + fn=gallery_index_changed, + _js="(a, b, c, d) => [dataset_tag_editor_gl_dataset_images_selected_index(), b, c, d]", + inputs=[dataset_gallery.nb_hidden_image_index ,self.tb_edit_caption, self.cb_copy_caption_automatically, self.cb_ask_save_when_caption_changed], + outputs=[self.nb_hidden_image_index_save_or_not] + [self.tb_caption, self.tb_edit_caption] + [self.tb_hidden_edit_caption] + ) + + def dialog_selected_save_caption_change(prev_idx:int, edit_caption: str, sort: bool = False): + return change_selected_image_caption(edit_caption, int(prev_idx), sort) + + self.btn_hidden_save_caption.click( + fn=dialog_selected_save_caption_change, + inputs=[self.nb_hidden_image_index_save_or_not, self.tb_hidden_edit_caption, self.cb_sort_caption_on_save], + outputs=o_update_filter_and_gallery + ) + + self.btn_copy_caption.click( + fn=lambda a:a, + inputs=[self.tb_caption], + outputs=[self.tb_edit_caption] + ) + + self.btn_append_caption.click( + fn=lambda a, b : b + (', ' if a and b else '') + a, + inputs=[self.tb_caption, self.tb_edit_caption], + outputs=[self.tb_edit_caption] + ) + + self.btn_prepend_caption.click( + fn=lambda a, b : a + (', ' if a and b else '') + b, + inputs=[self.tb_caption, self.tb_edit_caption], + outputs=[self.tb_edit_caption] + ) + + def interrogate_selected_image(interrogator_name: str, use_threshold_booru: bool, threshold_booru: float, use_threshold_waifu: bool, threshold_waifu: float): + if not interrogator_name: + return '' + threshold_booru = threshold_booru if use_threshold_booru else shared.opts.interrogate_deepbooru_score_threshold + threshold_waifu = threshold_waifu if use_threshold_waifu else -1 + return dte_module.interrogate_image(dataset_gallery.selected_path, interrogator_name, threshold_booru, threshold_waifu) + + self.btn_interrogate_si.click( + fn=interrogate_selected_image, + inputs=[self.dd_intterogator_names_si, load_dataset.cb_use_custom_threshold_booru, load_dataset.sl_custom_threshold_booru, load_dataset.cb_use_custom_threshold_waifu, load_dataset.sl_custom_threshold_waifu], + outputs=[self.tb_interrogate_selected_image] + ) + + self.btn_copy_interrogate.click( + fn=lambda a:a, + inputs=[self.tb_interrogate_selected_image], + outputs=[self.tb_edit_caption] + ) + + self.btn_append_interrogate.click( + fn=lambda a, b : b + (', ' if a and b else '') + a, + inputs=[self.tb_interrogate_selected_image, self.tb_edit_caption], + outputs=[self.tb_edit_caption] + ) + + self.btn_prepend_interrogate.click( + fn=lambda a, b : a + (', ' if a and b else '') + b, + inputs=[self.tb_interrogate_selected_image, self.tb_edit_caption], + outputs=[self.tb_edit_caption] + ) + + def change_selected_image_caption(tags_text: str, prev_idx:int, sort: bool = False): + idx = prev_idx + img_paths = dte_instance.get_filtered_imgpaths(filters=get_filters()) + + edited_tags = [t.strip() for t in tags_text.split(',')] + edited_tags = [t for t in edited_tags if t] + + if sort: + edited_tags = dte_instance.sort_tags(edited_tags) + + if 0 <= idx and idx < len(img_paths): + dte_instance.set_tags_by_image_path(imgpath=img_paths[idx], tags=edited_tags) + return update_filter_and_gallery() + + self.btn_apply_changes_selected_image.click( + fn=change_selected_image_caption, + inputs=[self.tb_edit_caption, self.cb_sort_caption_on_save], + outputs=o_update_filter_and_gallery + ) + self.btn_apply_changes_selected_image.click( + fn=lambda a:a, + inputs=[self.tb_edit_caption], + outputs=[self.tb_caption] + ) + self.btn_apply_changes_selected_image.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + diff --git a/scripts/ui/tab_filter_by_selection.py b/scripts/ui/tab_filter_by_selection.py new file mode 100644 index 0000000..57898c5 --- /dev/null +++ b/scripts/ui/tab_filter_by_selection.py @@ -0,0 +1,156 @@ +from typing import List +import gradio as gr + +from .ui_common import * +from .uibase import UIBase + +filters = dte_module.filters + + +class FilterBySelectionUI(UIBase): + def __init__(self): + self.path_filter = filters.PathFilter() + self.selected_index = -1 + self.selected_path = '' + self.tmp_selection = set() + + def get_current_txt_selection(self): + return f"""Selected Image : {self.selected_path}""" + + def create_ui(self, image_columns): + with gr.Row(visible=False): + self.btn_hidden_set_selection_index = gr.Button(elem_id="dataset_tag_editor_btn_hidden_set_selection_index") + self.nb_hidden_selection_image_index = gr.Number(value=-1) + gr.HTML("""Select images from the left gallery.""") + + with gr.Column(variant='panel'): + with gr.Row(): + self.btn_add_image_selection = gr.Button(value='Add selection [Enter]', elem_id='dataset_tag_editor_btn_add_image_selection') + self.btn_add_all_displayed_image_selection = gr.Button(value='Add ALL Displayed') + + self.gl_filter_images = gr.Gallery(label='Filter Images', elem_id="dataset_tag_editor_filter_gallery").style(grid=image_columns) + self.txt_selection = gr.HTML(value=self.get_current_txt_selection()) + + with gr.Row(): + self.btn_remove_image_selection = gr.Button(value='Remove selection [Delete]', elem_id='dataset_tag_editor_btn_remove_image_selection') + self.btn_invert_image_selection = gr.Button(value='Invert selection') + self.btn_clear_image_selection = gr.Button(value='Clear selection') + + self.btn_apply_image_selection_filter = gr.Button(value='Apply selection filter', variant='primary') + + def set_callbacks(self, o_update_filter_and_gallery, dataset_gallery, filter_by_tags, get_filters, update_filter_and_gallery): + def selection_index_changed(idx:int = -1): + idx = int(idx) if idx is not None else -1 + img_paths = arrange_selection_order(self.tmp_selection) + if idx < 0 or len(img_paths) <= idx: + self.selected_path = '' + idx = -1 + else: + self.selected_path = img_paths[idx] + self.selected_index = idx + return [self.get_current_txt_selection(), idx] + + self.btn_hidden_set_selection_index.click( + fn=selection_index_changed, + _js="(x) => [dataset_tag_editor_gl_filter_images_selected_index()]", + inputs=[self.nb_hidden_selection_image_index], + outputs=[self.txt_selection, self.nb_hidden_selection_image_index] + ) + + def add_image_selection(): + img_path = dataset_gallery.selected_path + if img_path: + self.tmp_selection.add(img_path) + return [dte_instance.images[p] for p in arrange_selection_order(self.tmp_selection)] + + self.btn_add_image_selection.click( + fn=add_image_selection, + outputs=[self.gl_filter_images] + ) + + def add_all_displayed_image_selection(): + img_paths = dte_instance.get_filtered_imgpaths(filters=get_filters()) + self.tmp_selection |= set(img_paths) + return [dte_instance.images[p] for p in arrange_selection_order(self.tmp_selection)] + + self.btn_add_all_displayed_image_selection.click( + fn=add_all_displayed_image_selection, + outputs=self.gl_filter_images + ) + + def invert_image_selection(): + img_paths = dte_instance.get_img_path_set() + self.tmp_selection = img_paths - self.tmp_selection + return [dte_instance.images[p] for p in arrange_selection_order(self.tmp_selection)] + + self.btn_invert_image_selection.click( + fn=invert_image_selection, + outputs=self.gl_filter_images + ) + + def remove_image_selection(): + img_path = self.selected_path + if img_path: + self.tmp_selection.remove(img_path) + self.selected_path = '' + self.selected_index = -1 + + return [ + [dte_instance.images[p] for p in arrange_selection_order(self.tmp_selection)], + self.get_current_txt_selection(), + -1 + ] + + self.btn_remove_image_selection.click( + fn=remove_image_selection, + outputs=[self.gl_filter_images, self.txt_selection, self.nb_hidden_selection_image_index] + ) + + def clear_image_selection(): + self.tmp_selection.clear() + self.selected_path = '' + self.selected_index = -1 + return[ + [], + self.get_current_txt_selection(), + -1 + ] + + self.btn_clear_image_selection.click( + fn=clear_image_selection, + outputs= + [self.gl_filter_images, self.txt_selection, self.nb_hidden_selection_image_index] + ) + + def clear_image_filter(): + self.path_filter = filters.PathFilter() + return clear_image_selection() + update_filter_and_gallery() + + filter_by_tags.btn_clear_all_filters.click( + fn=clear_image_filter, + outputs= + [self.gl_filter_images, self.txt_selection, self.nb_hidden_selection_image_index] + + o_update_filter_and_gallery + ) + + def apply_image_selection_filter(): + if len(self.tmp_selection) > 0: + self.path_filter = filters.PathFilter(self.tmp_selection, filters.PathFilter.Mode.INCLUSIVE) + else: + self.path_filter = filters.PathFilter() + return update_filter_and_gallery() + + self.btn_apply_image_selection_filter.click( + fn=apply_image_selection_filter, + outputs=o_update_filter_and_gallery + ) + self.btn_apply_image_selection_filter.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + + +def arrange_selection_order(paths: List[str]): + return sorted(paths) + + \ No newline at end of file diff --git a/scripts/ui/tab_filter_by_tags.py b/scripts/ui/tab_filter_by_tags.py new file mode 100644 index 0000000..c0466b4 --- /dev/null +++ b/scripts/ui/tab_filter_by_tags.py @@ -0,0 +1,87 @@ +import gradio as gr + +from .ui_common import * +from .uibase import UIBase +from .block_tag_filter import TagFilterUI + +filters = dte_module.filters + +class FilterByTagsUI(UIBase): + def __init__(self): + self.tag_filter_ui = TagFilterUI(tag_filter_mode=filters.TagFilter.Mode.INCLUSIVE) + self.tag_filter_ui_neg = TagFilterUI(tag_filter_mode=filters.TagFilter.Mode.EXCLUSIVE) + + def create_ui(self, cfg_filter_p, cfg_filter_n, get_filters): + self.cbg_hidden_dataset_filter = gr.CheckboxGroup(label='Dataset Filter', visible=False) + self.nb_hidden_dataset_filter_apply = gr.Number(label='Filter Apply', value=-1, visible=False) + + with gr.Row(): + self.btn_clear_tag_filters = gr.Button(value='Clear tag filters') + self.btn_clear_all_filters = gr.Button(value='Clear ALL filters') + + with gr.Tab(label='Positive Filter'): + with gr.Column(variant='panel'): + gr.HTML(value='Search tags / Filter images by tags (INCLUSIVE)') + logic_p = filters.TagFilter.Logic.OR if cfg_filter_p.logic=='OR' else filters.TagFilter.Logic.NONE if cfg_filter_p.logic=='NONE' else filters.TagFilter.Logic.AND + self.tag_filter_ui.create_ui(get_filters, logic_p, cfg_filter_p.sort_by, cfg_filter_p.sort_order, cfg_filter_p.sw_prefix, cfg_filter_p.sw_suffix, cfg_filter_p.sw_regex) + + with gr.Tab(label='Negative Filter'): + with gr.Column(variant='panel'): + gr.HTML(value='Search tags / Filter images by tags (EXCLUSIVE)') + logic_n = filters.TagFilter.Logic.AND if cfg_filter_n.logic=='AND' else filters.TagFilter.Logic.NONE if cfg_filter_n.logic=='NONE' else filters.TagFilter.Logic.OR + self.tag_filter_ui_neg.create_ui(get_filters, logic_n, cfg_filter_n.sort_by, cfg_filter_n.sort_order, cfg_filter_n.sw_prefix, cfg_filter_n.sw_suffix, cfg_filter_n.sw_regex) + + def set_callbacks(self, o_update_gallery, o_update_filter_and_gallery, batch_edit_captions, move_or_delete_files, update_gallery, update_filter_and_gallery, get_filters): + common_callback = lambda : \ + update_gallery() + \ + batch_edit_captions.get_common_tags(get_filters, self) + \ + [move_or_delete_files.get_current_move_or_delete_target_num()] + \ + [batch_edit_captions.tag_select_ui_remove.cbg_tags_update()] + + common_callback_output = \ + o_update_gallery + \ + [batch_edit_captions.tb_common_tags, batch_edit_captions.tb_edit_tags] + \ + [move_or_delete_files.ta_move_or_delete_target_dataset_num]+ \ + [batch_edit_captions.tag_select_ui_remove.cbg_tags] + + + self.tag_filter_ui.on_filter_update( + fn=lambda : + common_callback() + + [', '.join(self.tag_filter_ui.filter.tags)], + inputs=None, + outputs=common_callback_output + [batch_edit_captions.tb_sr_selected_tags], + _js='(...args) => {dataset_tag_editor_gl_dataset_images_close(); return args}' + ) + + self.tag_filter_ui_neg.on_filter_update( + fn=common_callback, + inputs=None, + outputs=common_callback_output, + _js='(...args) => {dataset_tag_editor_gl_dataset_images_close(); return args}' + ) + + self.tag_filter_ui.set_callbacks() + self.tag_filter_ui_neg.set_callbacks() + + self.btn_clear_tag_filters.click( + fn=lambda:self.clear_filters(update_filter_and_gallery), + outputs=o_update_filter_and_gallery + ) + + self.btn_clear_all_filters.click( + fn=lambda:self.clear_filters(update_filter_and_gallery), + outputs=o_update_filter_and_gallery + ) + + self.nb_hidden_dataset_filter_apply.change( + fn=lambda a, b: [a, b], + _js='(x, y) => [y>=0 ? dataset_tag_editor_gl_dataset_images_filter(x) : x, -1]', + inputs=[self.cbg_hidden_dataset_filter, self.nb_hidden_dataset_filter_apply], + outputs=[self.cbg_hidden_dataset_filter, self.nb_hidden_dataset_filter_apply] + ) + + def clear_filters(self, update_filter_and_gallery): + self.tag_filter_ui.clear_filter() + self.tag_filter_ui_neg.clear_filter() + return update_filter_and_gallery() diff --git a/scripts/ui/tab_move_or_delete_files.py b/scripts/ui/tab_move_or_delete_files.py new file mode 100644 index 0000000..a27b62d --- /dev/null +++ b/scripts/ui/tab_move_or_delete_files.py @@ -0,0 +1,110 @@ +from typing import List + +import gradio as gr + +from .ui_common import * +from .uibase import UIBase + +class MoveOrDeleteFilesUI(UIBase): + def __init__(self): + self.target_data = 'Selected One' + self.current_target_txt = '' + + def create_ui(self, cfg_file_move_delete): + gr.HTML(value='Note: Moved or deleted images will be unloaded.') + self.target_data = cfg_file_move_delete.range + self.rb_move_or_delete_target_data = gr.Radio(choices=['Selected One', 'All Displayed Ones'], value=cfg_file_move_delete.range, label='Move or Delete') + self.cbg_move_or_delete_target_file = gr.CheckboxGroup(choices=['Image File', 'Caption Text File', 'Caption Backup File'], label='Target', value=cfg_file_move_delete.target) + self.tb_move_or_delete_caption_ext = gr.Textbox(label='Caption File Ext', placeholder='txt', value=cfg_file_move_delete.caption_ext) + self.ta_move_or_delete_target_dataset_num = gr.HTML(value='Target dataset num: 0') + self.tb_move_or_delete_destination_dir = gr.Textbox(label='Destination Directory', value=cfg_file_move_delete.destination) + self.btn_move_or_delete_move_files = gr.Button(value='Move File(s)', variant='primary') + gr.HTML(value='Note: DELETE cannot be undone. The files will be deleted completely.') + self.btn_move_or_delete_delete_files = gr.Button(value='DELETE File(s)', variant='primary') + + def get_current_move_or_delete_target_num(self): + return self.current_target_txt + + def set_callbacks(self, o_update_filter_and_gallery, dataset_gallery, filter_by_tags, batch_edit_captions, filter_by_selection, edit_caption_of_selected_image, get_filters, update_filter_and_gallery): + def _get_current_move_or_delete_target_num(): + if self.target_data == 'Selected One': + self.current_target_txt = f'Target dataset num: {1 if dataset_gallery.selected_index != -1 else 0}' + elif self.target_data == 'All Displayed Ones': + img_paths = dte_instance.get_filtered_imgpaths(filters=get_filters()) + self.current_target_txt = f'Target dataset num: {len(img_paths)}' + else: + self.current_target_txt = f'Target dataset num: 0' + return self.current_target_txt + + update_args = { + 'fn' : _get_current_move_or_delete_target_num, + 'inputs' : None, + 'outputs' : [self.ta_move_or_delete_target_dataset_num] + } + + batch_edit_captions.btn_apply_edit_tags.click(**update_args) + + batch_edit_captions.btn_apply_sr_tags.click(**update_args) + + filter_by_selection.btn_apply_image_selection_filter.click(**update_args) + + filter_by_tags.btn_clear_tag_filters.click(**update_args) + + filter_by_tags.btn_clear_all_filters.click(**update_args) + + edit_caption_of_selected_image.btn_apply_changes_selected_image.click(**update_args) + + self.rb_move_or_delete_target_data.change(**update_args) + + def move_files(target_data: str, target_file: List[str], caption_ext: str, dest_dir: str): + move_img = 'Image File' in target_file + move_txt = 'Caption Text File' in target_file + move_bak = 'Caption Backup File' in target_file + if target_data == 'Selected One': + img_path = dataset_gallery.selected_path + if img_path: + dte_instance.move_dataset_file(img_path, caption_ext, dest_dir, move_img, move_txt, move_bak) + dte_instance.construct_tag_counts() + + elif target_data == 'All Displayed Ones': + dte_instance.move_dataset(dest_dir, caption_ext, get_filters(), move_img, move_txt, move_bak) + + return update_filter_and_gallery() + + self.btn_move_or_delete_move_files.click( + fn=move_files, + inputs=[self.rb_move_or_delete_target_data, self.cbg_move_or_delete_target_file, self.tb_move_or_delete_caption_ext, self.tb_move_or_delete_destination_dir], + outputs=o_update_filter_and_gallery + ) + self.btn_move_or_delete_move_files.click(**update_args) + self.btn_move_or_delete_move_files.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + + def delete_files(target_data: str, target_file: List[str], caption_ext: str): + delete_img = 'Image File' in target_file + delete_txt = 'Caption Text File' in target_file + delete_bak = 'Caption Backup File' in target_file + if target_data == 'Selected One': + img_path = dataset_gallery.selected_path + if img_path: + dte_instance.delete_dataset_file(img_path, delete_img, caption_ext, delete_txt, delete_bak) + dte_instance.construct_tag_counts() + + elif target_data == 'All Displayed Ones': + dte_instance.delete_dataset(caption_ext, get_filters(), delete_img, delete_txt, delete_bak) + + return update_filter_and_gallery() + + self.btn_move_or_delete_delete_files.click( + fn=delete_files, + inputs=[self.rb_move_or_delete_target_data, self.cbg_move_or_delete_target_file, self.tb_move_or_delete_caption_ext], + outputs=o_update_filter_and_gallery + ) + self.btn_move_or_delete_delete_files.click(**update_args) + self.btn_move_or_delete_delete_files.click( + fn=None, + _js='() => dataset_tag_editor_gl_dataset_images_close()' + ) + \ No newline at end of file diff --git a/scripts/ui/ui_classes.py b/scripts/ui/ui_classes.py new file mode 100644 index 0000000..024cb6e --- /dev/null +++ b/scripts/ui/ui_classes.py @@ -0,0 +1,15 @@ +from .block_toprow import ToprowUI +from .block_load_dataset import LoadDatasetUI +from .block_dataset_gallery import DatasetGalleryUI +from .block_gallery_state import GalleryStateUI +from .block_tag_filter import TagFilterUI +from .block_tag_select import TagSelectUI +from .tab_filter_by_tags import FilterByTagsUI +from .tab_filter_by_selection import FilterBySelectionUI +from .tab_batch_edit_captions import BatchEditCaptionsUI +from .tab_edit_caption_of_selected_image import EditCaptionOfSelectedImageUI +from .tab_move_or_delete_files import MoveOrDeleteFilesUI + +__all__ = [ + 'ToprowUI', 'LoadDatasetUI', 'DatasetGalleryUI', 'GalleryStateUI', 'TagFilterUI', 'TagSelectUI', 'FilterByTagsUI', 'FilterBySelectionUI', 'BatchEditCaptionsUI', 'EditCaptionOfSelectedImageUI', 'MoveOrDeleteFilesUI' +] \ No newline at end of file diff --git a/scripts/ui/ui_common.py b/scripts/ui/ui_common.py new file mode 100644 index 0000000..2b8dd51 --- /dev/null +++ b/scripts/ui/ui_common.py @@ -0,0 +1 @@ +from scripts.dte_instance import dte_instance, dte_module \ No newline at end of file diff --git a/scripts/ui/ui_instance.py b/scripts/ui/ui_instance.py new file mode 100644 index 0000000..6660347 --- /dev/null +++ b/scripts/ui/ui_instance.py @@ -0,0 +1,15 @@ +from .ui_classes import * + +__all__ = [ + 'toprow', 'load_dataset', 'dataset_gallery', 'gallery_state', 'filter_by_tags', 'filter_by_selection', 'batch_edit_captions', 'edit_caption_of_selected_image', 'move_or_delete_files' +] + +toprow = ToprowUI.get_instance() +load_dataset = LoadDatasetUI.get_instance() +dataset_gallery = DatasetGalleryUI.get_instance() +gallery_state = GalleryStateUI.get_instance() +filter_by_tags = FilterByTagsUI.get_instance() +filter_by_selection = FilterBySelectionUI.get_instance() +batch_edit_captions = BatchEditCaptionsUI.get_instance() +edit_caption_of_selected_image = EditCaptionOfSelectedImageUI.get_instance() +move_or_delete_files = MoveOrDeleteFilesUI.get_instance() diff --git a/scripts/ui/uibase.py b/scripts/ui/uibase.py new file mode 100644 index 0000000..d0274df --- /dev/null +++ b/scripts/ui/uibase.py @@ -0,0 +1,19 @@ +class Singleton(object): + @classmethod + def get_instance(cls): + if not hasattr(cls, "_instance"): + cls._instance = cls() + return cls._instance + +class UIBase(Singleton): + def create_ui(self, *args, **kwargs): + raise NotImplementedError() + def set_callbacks(self, *args, **kwargs): + raise NotImplementedError() + def func_to_set_value(self, name, type=None): + def func(value): + if type is not None: + value = type(value) + setattr(self, name, value) + return value + return func \ No newline at end of file