1580158
/
stable-diffusion-webui-dataset-tag-editor
mirror of https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

interrogate features #5 #6

pull/8/head
toshiaki1729 2022-11-12 13:07:46 +09:00
parent 605f8c1640
commit 89248bf87e
2 changed files with 32 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
dataset_tag_editor.py
View File

@ -17,6 +17,8 @@ class InterrogateMethod(Enum):
NONE = 0
PREFILL = 1
OVERWRITE = 2
PREPEND = 3
APPEND = 4
def interrogate_image_clip(path):
@ -176,7 +178,7 @@ class DatasetTagEditor:
def load_dataset(self, img_dir: str, recursive: bool = False, load_caption_from_filename: bool = True, interrogate_method: InterrogateMethod = InterrogateMethod.NONE, use_clip: bool = True):
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()
print(f'Loading dataset from {img_dir}')
try:
@ -188,8 +190,8 @@ class DatasetTagEditor:
self.dataset_dir = img_dir
if not (use_clip or cmd_opts.deepdanbooru):
print('Cannot interrogate without --deepdanbooru commandline option.')
if use_booru and not cmd_opts.deepdanbooru:
print('Cannot use DeepDanbooru without --deepdanbooru commandline option.')
print(f'Total {len(filepath_set)} files under the directory including not image files.')
@ -214,16 +216,26 @@ class DatasetTagEditor:
tokens = self.re_word.findall(caption_text)
caption_text = (shared.opts.dataset_filename_join_string or "").join(tokens)
if interrogate_method == InterrogateMethod.OVERWRITE or (not caption_text and interrogate_method == InterrogateMethod.PREFILL):
if interrogate_method != InterrogateMethod.NONE and ((interrogate_method != InterrogateMethod.PREFILL) or (interrogate_method == InterrogateMethod.PREFILL and not caption_text)):
try:
img = Image.open(img_path).convert('RGB')
except Exception as e:
print(e)
else:
interrogate_text = ''
if use_clip:
caption_text = shared.interrogator.generate_caption(img)
elif cmd_opts.deepdanbooru:
caption_text = deepbooru.get_tags_from_process(img)
interrogate_text += shared.interrogator.generate_caption(img)
if use_booru and cmd_opts.deepdanbooru:
tmp = deepbooru.get_tags_from_process(img)
interrogate_text += (', ' if interrogate_text and tmp else '') + tmp
if interrogate_method == InterrogateMethod.OVERWRITE:
caption_text = interrogate_text
elif interrogate_method == InterrogateMethod.PREPEND:
caption_text = interrogate_text + (', ' if interrogate_text and caption_text else '') + caption_text
else:
caption_text += (', ' if interrogate_text and caption_text else '') + interrogate_text
self.set_tags_by_image_path(img_path, [t.strip() for t in caption_text.split(',')])
@ -231,7 +243,7 @@ class DatasetTagEditor:
if interrogate_method != InterrogateMethod.NONE:
if use_clip:
shared.interrogator.load()
elif cmd_opts.deepdanbooru:
if use_booru and cmd_opts.deepdanbooru:
db_opts = deepbooru.create_deepbooru_opts()
db_opts[deepbooru.OPT_INCLUDE_RANKS] = False
deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts)
@ -242,7 +254,7 @@ class DatasetTagEditor:
if interrogate_method != InterrogateMethod.NONE:
if use_clip:
shared.interrogator.send_blip_to_ram()
elif cmd_opts.deepdanbooru:
if use_booru and cmd_opts.deepdanbooru:
deepbooru.release_process()
self.construct_tag_counts()

16
scripts/dataset_tag_editor_ui.py
View File

@ -38,7 +38,7 @@ def get_current_txt_selection():
return f"""Selected Image : {selection_selected_image_path}"""
def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool, load_caption_from_filename: bool, use_interrogator: str, use_clip: bool):
def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool, load_caption_from_filename: bool, use_interrogator: str, use_clip: bool, use_booru: bool):
global total_image_num, displayed_image_num, current_tag_filter, current_selection, tmp_selection_img_path_set, selected_image_path, selection_selected_image_path
interrogate_method = InterrogateMethod.NONE
@ -46,8 +46,12 @@ def load_files_from_dir(dir: str, sort_by: str, sort_order: str, recursive: bool
interrogate_method = InterrogateMethod.PREFILL
elif use_interrogator == 'Overwrite':
interrogate_method = InterrogateMethod.OVERWRITE
elif use_interrogator == 'Prepend':
interrogate_method = InterrogateMethod.PREPEND
elif use_interrogator == 'Append':
interrogate_method = InterrogateMethod.APPEND
dataset_tag_editor.load_dataset(img_dir=dir, recursive=recursive, load_caption_from_filename=load_caption_from_filename, interrogate_method=interrogate_method, use_clip=use_clip)
dataset_tag_editor.load_dataset(img_dir=dir, recursive=recursive, load_caption_from_filename=load_caption_from_filename, interrogate_method=interrogate_method, use_clip=use_clip, use_booru=use_booru)
img_paths, tags = dataset_tag_editor.get_filtered_imgpath_and_tags()
tags = arrange_tag_order(tags=tags, sort_by=sort_by, sort_order=sort_order)
total_image_num = displayed_image_num = len(dataset_tag_editor.get_img_path_set())
@ -283,8 +287,10 @@ def on_ui_tabs():
cb_load_recursive = gr.Checkbox(value=False, label='Load from subdirectories')
cb_load_caption_from_filename = gr.Checkbox(value=True, label='Load caption from filename if no text file')
with gr.Column():
rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite'], value='No', label='Use Interrogator Caption')
cb_use_clip_to_prefill = gr.Checkbox(value=False, label='Use BLIP interrogator')
rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite', 'Prepend', 'Append'], value='No', label='Use Interrogator Caption')
with gr.Row():
cb_use_clip_to_prefill = gr.Checkbox(value=False, label='Use BLIP')
cb_use_booru_to_prefill = gr.Checkbox(value=False, label='Use DeepDanbooru')
gl_dataset_images = gr.Gallery(label='Dataset Images', elem_id="dataset_tag_editor_dataset_gallery").style(grid=opts.dataset_editor_image_columns)
txt_filter = gr.HTML(value=get_current_txt_filter())
@ -381,7 +387,7 @@ def on_ui_tabs():
btn_load_datasets.click(
fn=load_files_from_dir,
inputs=[tb_img_directory, rd_sort_by, rd_sort_order, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill],
inputs=[tb_img_directory, rd_sort_by, rd_sort_order, cb_load_recursive, cb_load_caption_from_filename, rb_use_interrogator, cb_use_clip_to_prefill, cb_use_booru_to_prefill],
outputs=[gl_dataset_images, gl_selected_images, cbg_tags, tb_search_tags, txt_filter, txt_selection]
)
btn_load_datasets.click(