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]