diff --git a/scripts/dataset_tag_editor/dataset_tag_editor.py b/scripts/dataset_tag_editor/dataset_tag_editor.py index 331d2d6..8e99093 100644 --- a/scripts/dataset_tag_editor/dataset_tag_editor.py +++ b/scripts/dataset_tag_editor/dataset_tag_editor.py @@ -242,6 +242,105 @@ class DatasetTagEditor: return {k for k in self.dataset.datas.keys() if k} + def delete_dataset(self, filters: List[Dataset.Filter], delete_image: bool = False, delete_caption: bool = False, delete_backup: bool = False): + filtered_set = self.dataset.copy() + for filter in filters: + filtered_set.filter(filter) + for path in filtered_set.datas.keys(): + self.delete_dataset_file(path, delete_image, delete_caption, delete_backup) + + if delete_image: + self.dataset.remove(filtered_set) + print(list(self.dataset.datas.values())) + self.construct_tag_counts() + + + def move_dataset(self, dest_dir: str, filters: List[Dataset.Filter], move_image: bool = False, move_caption: bool = False, move_backup: bool = False): + filtered_set = self.dataset.copy() + for filter in filters: + filtered_set.filter(filter) + for path in filtered_set.datas.keys(): + self.move_dataset_file(path, dest_dir, move_image, move_caption, move_backup) + + if move_image: + self.construct_tag_counts() + + + def delete_dataset_file(self, img_path: str, delete_image: bool = False, delete_caption: bool = False, delete_backup: bool = False): + if img_path not in self.dataset.datas.keys(): + return + + if delete_image: + try: + if os.path.exists(img_path) and os.path.isfile(img_path): + os.remove(img_path) + self.dataset.remove_by_path(img_path) + print(f'Deleted {img_path}') + except Exception as e: + print(e) + + if delete_caption: + try: + filepath_noext, _ = os.path.splitext(img_path) + txt_path = filepath_noext + '.txt' + if os.path.exists(txt_path) and os.path.isfile(txt_path): + os.remove(txt_path) + print(f'Deleted {txt_path}') + except Exception as e: + print(e) + + if delete_backup: + try: + filepath_noext, _ = os.path.splitext(img_path) + for extnum in range(1000): + bak_path = filepath_noext + f'.{extnum:0>3d}' + if os.path.exists(bak_path) and os.path.isfile(bak_path): + os.remove(bak_path) + print(f'Deleted {bak_path}') + except Exception as e: + print(e) + + + def move_dataset_file(self, img_path: str, dest_dir: str, move_image: bool = False, move_caption: bool = False, move_backup: bool = False): + if img_path not in self.dataset.datas.keys(): + return + + if (move_image or move_caption or move_backup) and not os.path.exists(dest_dir): + os.mkdir(dest_dir) + + if move_image: + try: + dst_path = os.path.join(dest_dir, os.path.basename(img_path)) + if os.path.exists(img_path) and os.path.isfile(img_path): + os.replace(img_path, dst_path) + self.dataset.remove_by_path(img_path) + print(f'Moved {img_path} -> {dst_path}') + except Exception as e: + print(e) + + if move_caption: + try: + filepath_noext, _ = os.path.splitext(img_path) + txt_path = filepath_noext + '.txt' + dst_path = os.path.join(dest_dir, os.path.basename(txt_path)) + if os.path.exists(txt_path) and os.path.isfile(txt_path): + os.replace(txt_path, dst_path) + print(f'Moved {txt_path} -> {dst_path}') + except Exception as e: + print(e) + + if move_backup: + try: + filepath_noext, _ = os.path.splitext(img_path) + for extnum in range(1000): + bak_path = filepath_noext + f'.{extnum:0>3d}' + dst_path = os.path.join(dest_dir, os.path.basename(bak_path)) + if os.path.exists(bak_path) and os.path.isfile(bak_path): + os.replace(bak_path, dst_path) + print(f'Moved {bak_path} -> {dst_path}') + except Exception as e: + print(e) + def load_dataset(self, img_dir: str, recursive: bool = False, load_caption_from_filename: bool = True, interrogate_method: InterrogateMethod = InterrogateMethod.NONE, use_booru: bool = True, use_clip: bool = False): self.clear() diff --git a/scripts/main.py b/scripts/main.py index 102f992..9cbdd28 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -38,6 +38,16 @@ 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)}' + + def load_files_from_dir(dir: str, recursive: bool, load_caption_from_filename: bool, use_interrogator: str, use_clip: bool, use_booru: bool): global total_image_num, displayed_image_num, tmp_selection_img_path_set, gallery_selected_image_path, selection_selected_image_path, path_filter @@ -264,6 +274,44 @@ def search_and_replace(search_text: str, replace_text: str, target_text: str, us return update_filter_and_gallery() + update_common_tags() +# ================================================================ +# Callbacks for "Move or Delete Files" tab +# ================================================================ + +def move_files(target_data: str, target_file: List[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], 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, get_filters(), move_img, move_txt, move_bak) + + return update_filter_and_gallery() + update_common_tags() + + +def delete_files(target_data: str, target_file: List[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, delete_txt, delete_bak) + dataset_tag_editor.construct_tag_counts() + + elif target_data == 'All Displayed Ones': + dataset_tag_editor.delete_dataset(get_filters(), delete_img, delete_txt, delete_bak) + + return update_filter_and_gallery() + update_common_tags() + + # ================================================================ # Script Callbacks # ================================================================ @@ -399,6 +447,15 @@ def on_ui_tabs(): 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.""") + + 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'], label='Move or Delete') + cbg_move_or_delete_target_file = gr.CheckboxGroup(choices=['Image File', 'Caption Text File', 'Caption Backup File'], label='Target') + 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') + 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') #---------------------------------------------------------------- # Filter and Edit Tags tab @@ -408,8 +465,23 @@ def on_ui_tabs(): outputs=[txt_result] ) - tag_filter_ui.set_callbacks(on_filter_update=lambda:update_gallery() + update_common_tags() + [', '.join(tag_filter_ui.filter.tags)], outputs=[gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + [tb_sr_selected_tags]) - tag_filter_ui_neg.set_callbacks(on_filter_update=lambda:update_gallery() + update_common_tags() + [', '.join(tag_filter_ui.filter.tags)], outputs=[gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + [tb_sr_selected_tags]) + tag_filter_ui.set_callbacks( + on_filter_update=lambda: + update_gallery() + + update_common_tags() + + [', '.join(tag_filter_ui.filter.tags)] + + [get_current_move_or_delete_target_num(rb_move_or_delete_target_data, nb_hidden_image_index)], + outputs=[gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + [tb_sr_selected_tags] + [ta_move_or_delete_target_dataset_num] + ) + + tag_filter_ui_neg.set_callbacks( + on_filter_update=lambda: + update_gallery() + + update_common_tags() + + [', '.join(tag_filter_ui_neg.filter.tags)] + + [get_current_move_or_delete_target_num(rb_move_or_delete_target_data, nb_hidden_image_index)], + outputs=[gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + [tb_sr_selected_tags] + [ta_move_or_delete_target_dataset_num] + ) btn_load_datasets.click( fn=load_files_from_dir, @@ -501,6 +573,11 @@ def on_ui_tabs(): inputs=[nb_hidden_image_index], outputs=[tb_caption_selected_image, txt_gallery, nb_hidden_image_index] ) + btn_hidden_set_index.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_copy_caption.click( fn=lambda a:a, @@ -560,6 +637,27 @@ def on_ui_tabs(): outputs=[tb_caption_selected_image] ) + #---------------------------------------------------------------- + # 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_destination_dir, nb_hidden_image_index], + outputs=[tag_filter_ui.cbg_tags, tag_filter_ui_neg.cbg_tags, gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + ) + + btn_move_or_delete_delete_files.click( + fn=delete_files, + inputs=[rb_move_or_delete_target_data, cbg_move_or_delete_target_file, nb_hidden_image_index], + outputs=[tag_filter_ui.cbg_tags, tag_filter_ui_neg.cbg_tags, gl_dataset_images, nb_hidden_image_index, txt_gallery] + [tb_common_tags, tb_edit_tags] + ) + return [(dataset_tag_editor_interface, "Dataset Tag Editor", "dataset_tag_editor_interface")]