diff --git a/scripts/dataset_tag_editor/dataset_tag_editor.py b/scripts/dataset_tag_editor/dataset_tag_editor.py index 09c282f..0c64c76 100644 --- a/scripts/dataset_tag_editor/dataset_tag_editor.py +++ b/scripts/dataset_tag_editor/dataset_tag_editor.py @@ -272,6 +272,14 @@ class DatasetTagEditor: self.set_tags_by_image_path(path, res) + def remove_tags(self, tags: Set[str], filters: List[Dataset.Filter] = []): + img_paths = self.get_filtered_imgpaths(filters) + for path in img_paths: + res = self.dataset.get_data_tags(path) + res = [t for t in res if t not in tags] + self.set_tags_by_image_path(path, res) + + def get_img_path_list(self): return [k for k in self.dataset.datas.keys() if k] diff --git a/scripts/dataset_tag_editor/ui.py b/scripts/dataset_tag_editor/ui.py index d97fd62..ffb2656 100644 --- a/scripts/dataset_tag_editor/ui.py +++ b/scripts/dataset_tag_editor/ui.py @@ -86,4 +86,82 @@ class TagFilterUI: def clear_filter(self): self.filter = TagFilter({}, self.logic, self.filter_mode) self.filter_word = '' - self.selected_tags = [] \ No newline at end of file + self.selected_tags = [] + + +class TagSelectUI: + def __init__(self, dataset_tag_editor: DatasetTagEditor): + 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:[] + + + def create_ui(self, get_filters: Callable[[], List[Dataset.Filter]], sort_by = 'Alphabetical Order', sort_order = 'Ascending'): + self.get_filters = get_filters + + import gradio as gr + self.tb_search_tags = gr.Textbox(label='Search Tags', interactive=True) + with gr.Row(): + self.rb_sort_by = gr.Radio(choices=['Alphabetical Order', 'Frequency'], 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.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 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) + self.tags = set(self.dataset_tag_editor.get_filtered_tags(self.get_filters(), filter_tags=True)) + 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/main.py b/scripts/main.py index a2a273c..87c6f8b 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -8,13 +8,14 @@ from collections import namedtuple from scripts.dataset_tag_editor.dataset_tag_editor import DatasetTagEditor, InterrogateMethod from scripts.dataset_tag_editor.filters import TagFilter, PathFilter -from scripts.dataset_tag_editor.ui import TagFilterUI +from scripts.dataset_tag_editor.ui import TagFilterUI, TagSelectUI dataset_tag_editor = DatasetTagEditor() path_filter = PathFilter() tag_filter_ui:TagFilterUI = None tag_filter_ui_neg:TagFilterUI = None +tag_select_ui_remove:TagSelectUI = None def get_filters(): return [path_filter, tag_filter_ui.get_filter(), tag_filter_ui_neg.get_filter()] @@ -36,14 +37,14 @@ CONFIG_PATH = os.path.join(scripts.basedir(), 'config.json') GeneralConfig = namedtuple('GeneralConfig', ['backup', 'dataset_dir', 'load_recursive', 'load_caption_from_filename', 'use_interrogator', 'use_clip_to_prefill', 'use_booru_to_prefill', 'use_waifu_to_prefill']) FilterConfig = namedtuple('FilterConfig', ['sort_by', 'sort_order', 'logic']) -BatchEditConfig = namedtuple('BatchEditConfig', ['show_only_selected', 'prepend', 'use_regex', 'target']) +BatchEditConfig = namedtuple('BatchEditConfig', ['show_only_selected', 'prepend', 'use_regex', 'target', 'sory_by', 'sort_order']) EditSelectedConfig = namedtuple('EditSelectedConfig', ['auto_copy', 'warn_change_not_saved']) MoveDeleteConfig = namedtuple('MoveDeleteConfig', ['range', 'target', 'destination']) CFG_GENERAL_DEFAULT = GeneralConfig(True, '', False, True, 'No', False, 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') +CFG_BATCH_EDIT_DEFAULT = BatchEditConfig(True, False, False, 'Only Selected Tags', 'Alphabetical Order', 'Ascending') CFG_EDIT_SELECTED_DEFAULT = EditSelectedConfig(False, False) CFG_MOVE_DELETE_DEFAULT = MoveDeleteConfig('Selected One', [], '') @@ -83,8 +84,8 @@ def write_filter_config(sort_by_p: str, sort_order_p: str, logic_p: str, sort_by cfg_n = FilterConfig(sort_by_n, sort_order_n, logic_n) config.write({'positive':cfg_p._asdict(), 'negative':cfg_n._asdict()}, 'filter') -def write_batch_edit_config(show_only_selected: bool, prepend: bool, use_regex: bool, target: str): - cfg = BatchEditConfig(show_only_selected, prepend, use_regex, target) +def write_batch_edit_config(show_only_selected: bool, prepend: bool, use_regex: bool, target: str, sort_by: str, sort_order: str): + cfg = BatchEditConfig(show_only_selected, prepend, use_regex, target, sort_by, sort_order) config.write(cfg._asdict(), 'batch_edit') def write_edit_selected_config(auto_copy: bool, warn_change_not_saved: bool): @@ -185,7 +186,10 @@ def load_files_from_dir(dir: str, recursive: bool, load_caption_from_filename: b [], 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() + ] +\ + [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 update_gallery(): @@ -203,7 +207,12 @@ def update_gallery(): def update_filter_and_gallery(): - return [tag_filter_ui.cbg_tags_update(), tag_filter_ui_neg.cbg_tags_update()] + update_gallery() + update_common_tags() + [', '.join(tag_filter_ui.filter.tags)] + return \ + [tag_filter_ui.cbg_tags_update(), 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(): @@ -421,6 +430,11 @@ def remove_duplicated_tags(): 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() + + # ================================================================ # Callbacks for "Move or Delete Files" tab # ================================================================ @@ -464,7 +478,7 @@ def delete_files(target_data: str, target_file: List[str], idx: int = -1): # ================================================================ def on_ui_tabs(): - global tag_filter_ui, tag_filter_ui_neg + global tag_filter_ui, tag_filter_ui_neg, tag_select_ui_remove config.load() cfg_general = read_general_config() @@ -475,7 +489,7 @@ def on_ui_tabs(): tag_filter_ui = TagFilterUI(dataset_tag_editor, tag_filter_mode=TagFilter.Mode.INCLUSIVE) tag_filter_ui_neg = TagFilterUI(dataset_tag_editor, tag_filter_mode=TagFilter.Mode.EXCLUSIVE) - + tag_select_ui_remove = TagSelectUI(dataset_tag_editor) with gr.Blocks(analytics_enabled=False) as dataset_tag_editor_interface: with gr.Row(visible=False): @@ -569,41 +583,48 @@ def on_ui_tabs(): btn_apply_image_selection_filter = gr.Button(value='Apply selection filter', variant='primary') with gr.Tab(label='Batch Edit Captions'): - 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.Column(variant='panel'): - btn_remove_duplicate = gr.Button(value='Remove duplicated tags') + 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='Delete'): + with gr.Column(variant='panel'): + gr.HTML('Remove duplicated tags in the images displayed.') + btn_remove_duplicate = gr.Button(value='Remove duplicated 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) with gr.Tab(label='Edit Caption of Selected Image'): with gr.Tab(label='Read Caption from Selected Image'): @@ -639,7 +660,11 @@ def on_ui_tabs(): btn_move_or_delete_move_files = gr.Button(value='Move File(s)', variant='primary') 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] + 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] #---------------------------------------------------------------- # General @@ -647,7 +672,7 @@ def on_ui_tabs(): configurable_components = \ [cb_backup, tb_img_directory, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill, cb_use_booru_to_prefill, cb_use_waifu_to_prefill] +\ [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] +\ - [cb_show_only_tags_selected, cb_prepend_tags, cb_use_regex, rb_sr_replace_target] +\ + [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] +\ [cb_copy_caption_automatically, cb_ask_save_when_caption_changed] +\ [rb_move_or_delete_target_data, cbg_move_or_delete_target_file, tb_move_or_delete_destination_dir] @@ -664,9 +689,9 @@ def on_ui_tabs(): def save_settings_callback(*a): write_general_config(*a[:8]) write_filter_config(*a[8:14]) - write_batch_edit_config(*a[14:18]) - write_edit_selected_config(*a[18:20]) - write_move_delete_config(*a[20:]) + write_batch_edit_config(*a[14:20]) + write_edit_selected_config(*a[20:22]) + write_move_delete_config(*a[22:]) config.save() btn_save_setting_as_default.click( @@ -703,24 +728,38 @@ def on_ui_tabs(): update_gallery() + update_common_tags() + [', '.join(tag_filter_ui.filter.tags)] + - [get_current_move_or_delete_target_num(a, b)], + [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] + 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] ) 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)], + [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] + 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] ) btn_load_datasets.click( fn=load_files_from_dir, inputs=[tb_img_directory, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill, cb_use_booru_to_prefill, cb_use_waifu_to_prefill], - outputs=[gl_dataset_images, gl_filter_images, txt_gallery, txt_selection] + [cbg_hidden_dataset_filter, nb_hidden_dataset_filter_apply] + o_filter_and_gallery + 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], @@ -841,6 +880,13 @@ def on_ui_tabs(): outputs=o_filter_and_gallery ) + tag_select_ui_remove.set_callbacks() + + btn_remove_selected.click( + fn=remove_selected_tags, + outputs=o_filter_and_gallery + ) + #---------------------------------------------------------------- # Edit Caption of Selected Image tab