toshiaki1729
parent
3c5b7083a0
commit
0f9b290eb7
|
|
@ -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()
|
||||
|
|
|
|||
102
scripts/main.py
102
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='<b>Note: </b>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")]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue