From 4e00a6957044cdd3e1497b0c96fd4cd2642715f7 Mon Sep 17 00:00:00 2001 From: toshiaki1729 <116595002+toshiaki1729@users.noreply.github.com> Date: Thu, 9 Feb 2023 20:36:27 +0900 Subject: [PATCH] implement #37: prefix/suffix and regex on tag search --- .../dataset_tag_editor/dataset_tag_editor.py | 37 +++++++++- scripts/dataset_tag_editor/ui.py | 73 +++++++++++++++++-- scripts/main.py | 22 +++--- 3 files changed, 110 insertions(+), 22 deletions(-) diff --git a/scripts/dataset_tag_editor/dataset_tag_editor.py b/scripts/dataset_tag_editor/dataset_tag_editor.py index f41beee..330eaab 100644 --- a/scripts/dataset_tag_editor/dataset_tag_editor.py +++ b/scripts/dataset_tag_editor/dataset_tag_editor.py @@ -176,16 +176,45 @@ class DatasetTagEditor: return [self.img_idx.get(p) for p in img_paths] - def get_filtered_tags(self, filters: List[filters.Filter] = [], filter_word: str = '', filter_tags = True): + def get_filtered_tags(self, filters: List[filters.Filter] = [], filter_word: str = '', filter_tags = True, prefix=False, suffix=False, regex=False): if filter_tags: filtered_set = self.dataset.copy() for filter in filters: filtered_set.filter(filter) - tags = filtered_set.get_tagset() + tags:Set[str] = filtered_set.get_tagset() else: - tags = self.dataset.get_tagset() + tags:Set[str] = self.dataset.get_tagset() - return {tag for tag in tags if filter_word in tag} + 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 def cleanup_tags(self, tags: List[str]): diff --git a/scripts/dataset_tag_editor/ui.py b/scripts/dataset_tag_editor/ui.py index 3446d59..4e702e5 100644 --- a/scripts/dataset_tag_editor/ui.py +++ b/scripts/dataset_tag_editor/ui.py @@ -6,7 +6,7 @@ filters = dte.filters class TagFilterUI: - def __init__(self, dataset_tag_editor: dte.DatasetTagEditor, tag_filter_mode: filters.TagFilter.Mode = filters.TagFilter.Mode.INCLUSIVE): + 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' @@ -16,19 +16,29 @@ class TagFilterUI: 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'): + 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') @@ -39,17 +49,36 @@ class TagFilterUI: 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() @@ -76,9 +105,9 @@ class TagFilterUI: 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) + 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) + 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) @@ -99,7 +128,7 @@ class TagFilterUI: class TagSelectUI: - def __init__(self, dataset_tag_editor: dte.DatasetTagEditor): + def __init__(self, dataset_tag_editor): self.filter_word = '' self.sort_by = 'Alphabetical Order' self.sort_order = 'Ascending' @@ -107,13 +136,23 @@ class TagSelectUI: 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'): + 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') @@ -125,6 +164,9 @@ class TagSelectUI: 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) @@ -137,6 +179,21 @@ class TagSelectUI: 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() @@ -166,8 +223,8 @@ class TagSelectUI: def cbg_tags_update(self): - tags = self.dataset_tag_editor.get_filtered_tags(self.get_filters(), self.filter_word, True) - self.tags = set(self.dataset_tag_editor.get_filtered_tags(self.get_filters(), filter_tags=True)) + 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) diff --git a/scripts/main.py b/scripts/main.py index 5373ebf..f237f8a 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -56,15 +56,15 @@ GeneralConfig = namedtuple('GeneralConfig', [ 'meta_save_as_caption', 'meta_use_full_path' ]) -FilterConfig = namedtuple('FilterConfig', ['sort_by', 'sort_order', 'logic']) -BatchEditConfig = namedtuple('BatchEditConfig', ['show_only_selected', 'prepend', 'use_regex', 'target', 'sory_by', 'sort_order']) +FilterConfig = namedtuple('FilterConfig', ['sw_prefix', 'sw_suffix', 'sw_regex','sort_by', 'sort_order', 'logic']) +BatchEditConfig = namedtuple('BatchEditConfig', ['show_only_selected', 'prepend', 'use_regex', 'target', 'sw_prefix', 'sw_suffix', 'sw_regex', 'sory_by', 'sort_order']) EditSelectedConfig = namedtuple('EditSelectedConfig', ['auto_copy', 'warn_change_not_saved', 'use_interrogator_name']) MoveDeleteConfig = namedtuple('MoveDeleteConfig', ['range', 'target', 'caption_ext', 'destination']) CFG_GENERAL_DEFAULT = GeneralConfig(True, '', '.txt', False, True, 'No', [], False, 0.7, False, 0.5, False, '', '', True, False, False) -CFG_FILTER_P_DEFAULT = FilterConfig('Alphabetical Order', 'Ascending', 'AND') -CFG_FILTER_N_DEFAULT = FilterConfig('Alphabetical Order', 'Ascending', 'OR') -CFG_BATCH_EDIT_DEFAULT = BatchEditConfig(True, False, False, 'Only Selected Tags', 'Alphabetical Order', 'Ascending') +CFG_FILTER_P_DEFAULT = FilterConfig(False, False, False, 'Alphabetical Order', 'Ascending', 'AND') +CFG_FILTER_N_DEFAULT = FilterConfig(False, False, False, 'Alphabetical Order', 'Ascending', 'OR') +CFG_BATCH_EDIT_DEFAULT = BatchEditConfig(True, False, False, 'Only Selected Tags', False, False, False, 'Alphabetical Order', 'Ascending') CFG_EDIT_SELECTED_DEFAULT = EditSelectedConfig(False, False, '') CFG_MOVE_DELETE_DEFAULT = MoveDeleteConfig('Selected One', [], '.txt', '') @@ -614,13 +614,13 @@ def on_ui_tabs(): 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) + 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) + 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) with gr.Tab(label='Filter by Selection'): with gr.Row(visible=False): @@ -685,7 +685,7 @@ def on_ui_tabs(): 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) + 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='Edit Caption of Selected Image'): with gr.Tab(label='Read Caption from Selected Image'): @@ -737,8 +737,10 @@ def on_ui_tabs(): 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 ] - components_filter = [tag_filter_ui.rb_sort_by, tag_filter_ui.rb_sort_order, tag_filter_ui.rb_logic, 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.rb_sort_by, tag_select_ui_remove.rb_sort_order] + 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_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]