diff --git a/README.md b/README.md index 50a07c7..958047a 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,10 @@ You can also install it manually by running the following command from within th git clone https://github.com/AlUlkesh/stable-diffusion-webui-images-browser/ extensions/stable-diffusion-webui-images-browser +and restart your stable-diffusion-webui, then you can see the new tab "Image Browser". + ## Recent updates +- Additional sorting and filtering by EXIF data including .txt file information - Recyle bin option - Add/Remove from saved directories, via buttons - New dropdown with subdirs diff --git a/scripts/images_history.py b/scripts/images_history.py index c0b46bf..c3ecb9d 100644 --- a/scripts/images_history.py +++ b/scripts/images_history.py @@ -8,22 +8,31 @@ import random import gradio as gr import modules.extras import modules.ui +import json +import re from modules.shared import opts, cmd_opts from modules.ui_components import ToolButton -from modules import shared, scripts +from modules import shared, scripts, images +from modules.ui_common import plaintext_to_html from modules import script_callbacks from PIL import Image from pathlib import Path from send2trash import send2trash from typing import List, Tuple +from PIL.ExifTags import TAGS +from PIL.PngImagePlugin import PngImageFile, PngInfo + favorite_tab_name = "Favorites" -tabs_list = ["txt2img", "img2img", "txt2img-grids", "img2img-grids", "Extras", favorite_tab_name, "Others"] #txt2img-grids and img2img-grids added by HaylockGrant +tabs_list = ["txt2img", "img2img", "instruct-pix2pix", "txt2img-grids", "img2img-grids", "Extras", favorite_tab_name, "Others"] #txt2img-grids and img2img-grids added by HaylockGrant num_of_imgs_per_page = 0 loads_files_num = 0 path_recorder_filename = os.path.join(scripts.basedir(), "path_recorder.txt") path_recorder_filename_tmp = f"{path_recorder_filename}.tmp" image_ext_list = [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"] +cur_ranking_value="0" +finfo_aes = {} +finfo_exif = {} none_select = "Nothing selected" refresh_symbol = '\U0001f504' # 🔄 up_symbol = '\U000025b2' # ▲ @@ -144,6 +153,21 @@ def save_image(file_name): else: return "
Image not found (may have been already moved)
" +def create_ranked_file(filename, ranking): + ranking_file = 'ranking.json' + + if not os.path.isfile(ranking_file): + data = {} + + else: + with open(ranking_file, 'r') as file: + data = json.load(file) + + data[filename] = ranking + + with open(ranking_file, 'w') as file: + json.dump(data, file) + def delete_image(delete_num, name, filenames, image_index, visible_num): if name == "": return filenames, delete_num @@ -191,23 +215,134 @@ def traverse_all_files(curr_path, image_list, tabname_box, img_path_depth) -> Li current_depth = current_depth - 1 return image_list -def get_all_images(dir_name, sort_by, sort_order, keyword, tabname_box, img_path_depth): + +def cache_aes(fileinfos): + aes_cache_file = 'aes_scores.json' + aes_cache = {} + + if os.path.isfile(aes_cache_file): + with open(aes_cache_file, 'r') as file: + aes_cache = json.load(file) + + for fi_info in fileinfos: + if fi_info[0] in aes_cache: + finfo_aes[fi_info[0]] = aes_cache[fi_info[0]] + else: + finfo_aes[fi_info[0]] = "0" + aes_cache[fi_info[0]] = "0" + try: + image = PngImageFile(fi_info[0]) + allExif = modules.extras.run_pnginfo(image)[1] + if allExif and allExif != "": + m = re.search("aesthetic_score: (\d+\.\d+)", allExif) + if m: + finfo_aes[fi_info[0]] = m.group(1) + aes_cache[fi_info[0]] = m.group(1) + else: + try: + filename = os.path.splitext(fi_info[0])[0] + ".txt" + geninfo = "" + with open(filename) as f: + for line in f: + geninfo += line + finfo_exif[fi_info[0]] = geninfo + exif_cache[fi_info[0]] = geninfo + except Exception: + print(f"No exif for {fi_info[0]}") + except SyntaxError: + print(f"Non-PNG file in directory when doing AES check: {fi_info[0]}") + + with open(aes_cache_file, 'w') as file: + json.dump(aes_cache, file) + +def cache_exif(fileinfos): + exif_cache_file = 'exif_data.json' + exif_cache = {} + if os.path.isfile(exif_cache_file): + with open(exif_cache_file, 'r') as file: + exif_cache = json.load(file) + + for fi_info in fileinfos: + if fi_info[0] in exif_cache: + #print(f"{fi_info[0]} found in EXIF cache!") + finfo_exif[fi_info[0]] = exif_cache[fi_info[0]] + else: + #print(f"{fi_info[0]} NOT found in exif cache!") + finfo_exif[fi_info[0]] = "0" + try: + image = PngImageFile(fi_info[0]) + allExif = modules.extras.run_pnginfo(image)[1] + if allExif: + finfo_exif[fi_info[0]] = allExif + exif_cache[fi_info[0]] = allExif + else: + try: + filename = os.path.splitext(fi_info[0])[0] + ".txt" + geninfo = "" + with open(filename) as f: + for line in f: + geninfo += line + finfo_exif[fi_info[0]] = geninfo + exif_cache[fi_info[0]] = geninfo + except Exception: + print(f"No EXIF in PNG or txt file fpr {fi_info[0]}") + #print(f"{fi_info[0]} exif added: {allExif}!") + except SyntaxError: + print(f"Non-PNG file in directory when doing EXIF check: {fi_info[0]}") + + with open(exif_cache_file, 'w') as file: + json.dump(exif_cache, file) + +def atof(text): + try: + retval = float(text) + except ValueError: + retval = text + return retval + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + float regex comes from https://stackoverflow.com/a/12643073/190597 + ''' + return [ atof(c) for c in re.split(r'[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)', text) ] + + +def get_all_images(dir_name, sort_by, sort_order, keyword, ranking_filter, aes_filter, desc, exif_keyword, tabname_box, img_path_depth): global current_depth current_depth = 0 fileinfos = traverse_all_files(dir_name, [], tabname_box, img_path_depth) keyword = keyword.strip(" ") + + cache_aes(fileinfos) + cache_exif(fileinfos) + if len(keyword) != 0: fileinfos = [x for x in fileinfos if keyword.lower() in x[0].lower()] + filenames = [finfo[0] for finfo in fileinfos] + if len(exif_keyword) != 0: + fileinfos = [x for x in fileinfos if exif_keyword.lower() in finfo_exif[x[0]].lower()] + filenames = [finfo[0] for finfo in fileinfos] + if len(aes_filter) != 0: + fileinfos = [x for x in fileinfos if finfo_aes[x[0]] >= aes_filter] + filenames = [finfo[0] for finfo in fileinfos] + if ranking_filter != "All": + fileinfos = [x for x in fileinfos if get_ranking(x[0]) in ranking_filter] + filenames = [finfo[0] for finfo in fileinfos] if sort_by == "date": if sort_order == up_symbol: fileinfos = sorted(fileinfos, key=lambda x: x[1].st_mtime) else: fileinfos = sorted(fileinfos, key=lambda x: -x[1].st_mtime) + filenames = [finfo[0] for finfo in fileinfos] elif sort_by == "path name": if sort_order == up_symbol: fileinfos = sorted(fileinfos) else: fileinfos = sorted(fileinfos, reverse=True) + filenames = [finfo[0] for finfo in fileinfos] elif sort_by == "aesthetic_score": if sort_order == up_symbol: fileinfos = sorted(fileinfos, key=lambda x: get_image_aesthetic_score(x[0])) @@ -215,8 +350,42 @@ def get_all_images(dir_name, sort_by, sort_order, keyword, tabname_box, img_path fileinfos = sorted(fileinfos, key=lambda x: -get_image_aesthetic_score(x[0])) elif sort_by == "random": random.shuffle(fileinfos) - - filenames = [finfo[0] for finfo in fileinfos] + filenames = [finfo[0] for finfo in fileinfos] + elif sort_by == "ranking": + finfo_ranked = {} + for fi_info in fileinfos: + finfo_ranked[fi_info[0]] = get_ranking(fi_info[0]) + if not desc: + fileinfos = dict(sorted(finfo_ranked.items(), key=lambda x: (x[1], x[0]))) + else: + fileinfos = dict(reversed(sorted(finfo_ranked.items(), key=lambda x: (x[1], x[0])))) + filenames = [finfo for finfo in fileinfos] + elif sort_by == "aes": + fileinfo_aes = {} + for finfo in fileinfos: + fileinfo_aes[finfo[0]] = finfo_aes[finfo[0]] + if desc: + fileinfos = dict(reversed(sorted(fileinfo_aes.items(), key=lambda x: (x[1], x[0])))) + else: + fileinfos = dict(sorted(fileinfo_aes.items(), key=lambda x: (x[1], x[0]))) + filenames = [finfo for finfo in fileinfos] + else: + sort_values = {} + exif_info = dict(finfo_exif) + if exif_info: + for k, v in exif_info.items(): + match = re.search(r'(?<='+ sort_by.lower() + ":" ').*?(?=(,|$))', v.lower()) + if match: + sort_values[k] = match.group() + else: + sort_values[k] = "0" + if desc: + fileinfos = dict(reversed(sorted(fileinfos, key=lambda x: natural_keys(sort_values[x[0]])))) + else: + fileinfos = dict(sorted(fileinfos, key=lambda x: natural_keys(sort_values[x[0]]))) + filenames = [finfo for finfo in fileinfos] + else: + filenames = [finfo for finfo in fileinfos] return filenames def get_image_aesthetic_score(img_path): @@ -229,7 +398,7 @@ def get_image_aesthetic_score(img_path): except KeyError: return 0 -def get_image_page(img_path, page_index, filenames, keyword, sort_by, sort_order, tabname_box, img_path_depth): +def get_image_page(img_path, page_index, filenames, keyword, sort_by, ranking_filter, aes_filter, sort_order, exif_keyword): img_path, _ = pure_path(img_path) if not cmd_opts.administrator: head = os.path.abspath(".") @@ -238,7 +407,7 @@ def get_image_page(img_path, page_index, filenames, keyword, sort_by, sort_order warning = warning_permission.format(img_path) return None, 0, None, "", "", "", None, None, warning if page_index == 1 or page_index == 0 or len(filenames) == 0: - filenames = get_all_images(img_path, sort_by, sort_order, keyword, tabname_box, img_path_depth) + filenames = get_all_images(img_path, sort_by, sort_order, keyword, tabname_box, img_path_depth, ranking_filter, aes_filter, exif_keyword)) page_index = int(page_index) length = len(filenames) max_page_index = length // num_of_imgs_per_page + 1 @@ -256,11 +425,22 @@ def get_image_page(img_path, page_index, filenames, keyword, sort_by, sort_order load_info += "" return filenames, gr.update(value=page_index, label=f"Page Index (of {max_page_index} pages)"), image_list, "", "", "", visible_num, load_info +def get_current_file(tabname_box, num, page_index, filenames): + file = filenames[int(num) + int((page_index - 1) * num_of_imgs_per_page)] + return file + def show_image_info(tabname_box, num, page_index, filenames): - file = filenames[int(num) + int((page_index - 1) * num_of_imgs_per_page)] + file = filenames[int(num) + int((page_index - 1) * num_of_imgs_per_page)] tm = "
" + time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(os.path.getmtime(file))) + "
" return file, tm, num, file, "" +def show_next_image_info(tabname_box, num, page_index, filenames, auto_next): + file = filenames[int(num) + int((page_index - 1) * num_of_imgs_per_page)] + tm = "
" + time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(os.path.getmtime(file))) + "
" + if auto_next: + num = int(num) + 1 + return file, tm, num, file, "" + def change_dir(img_dir, path_recorder, load_switch, img_path_history, img_path_depth, img_path): warning = None img_path, _ = pure_path(img_path) @@ -296,14 +476,45 @@ def change_dir(img_dir, path_recorder, load_switch, img_path_history, img_path_d def update_move_text(unused): return f'{"Move" if not opts.images_copy_image else "Copy"} to favorites' +def get_ranking(filename): + ranking_file = 'ranking.json' + ranking_value = "None" + if os.path.isfile(ranking_file): + with open(ranking_file, 'r') as file: + data = json.load(file) + if filename in data: + ranking_value = data[filename] + + return ranking_value + +def get_ranking(filename): + ranking_file = 'ranking.json' + ranking_value = "None" + if os.path.isfile(ranking_file): + with open(ranking_file, 'r') as file: + data = json.load(file) + if filename in data: + ranking_value = data[filename] + + return ranking_value + def create_tab(tabname): custom_dir = False path_recorder = {} path_recorder_formatted = [] + + try: + if opts.outdir_ip2p_samples: + ip2p_dirname = opts.outdir_ip2p_samples + except AttributeError: + ip2p_dirname = "outputs/ip2p-images" + if tabname == "txt2img": dir_name = opts.outdir_txt2img_samples elif tabname == "img2img": dir_name = opts.outdir_img2img_samples + elif tabname == "instruct-pix2pix": + dir_name = ip2p_dirname elif tabname == "txt2img-grids": #added by HaylockGrant to add a new tab for grid images dir_name = opts.outdir_txt2img_grids elif tabname == "img2img-grids": #added by HaylockGrant to add a new tab for grid images @@ -346,7 +557,7 @@ def create_tab(tabname): with gr.Row(visible= not custom_dir, elem_id=tabname + "_images_history") as main_panel: with gr.Column(): with gr.Row(): - with gr.Column(scale=2): + with gr.Column(scale=2): with gr.Row(): first_page = gr.Button('First Page') prev_page = gr.Button('Prev Page') @@ -354,19 +565,28 @@ def create_tab(tabname): refresh_index_button = ToolButton(value=refresh_symbol) next_page = gr.Button('Next Page') end_page = gr.Button('End Page') + with gr.Column(scale=10): + ranking = gr.Radio(value="None", choices=["1", "2", "3", "4", "5", "None"], label="ranking", interactive="true") + auto_next = gr.Checkbox(label="Next Image After Ranking (To be implemented)", interactive="false") history_gallery = gr.Gallery(show_label=False, elem_id=tabname + "_images_history_gallery").style(grid=opts.images_history_page_columns) with gr.Row() as delete_panel: with gr.Column(scale=1): delete_num = gr.Number(value=1, interactive=True, label="delete next") with gr.Column(scale=3): delete = gr.Button('Delete', elem_id=tabname + "_images_history_del_button") + + with gr.Column(scale=1): + with gr.Row(scale=0.5): + sort_by = gr.Dropdown(value="date", choices=["path name", "date", "aesthetic_score", "random", "cfg scale", "steps", "seed", "sampler", "size", "model", "model hash", "ranking"], label="sort by") + sort_order = ToolButton(value=down_symbol) + with gr.Row(): + keyword = gr.Textbox(value="", label="filename keyword") + exif_keyword = gr.Textbox(value="", label="exif keyword") - with gr.Column(): + with gr.Column(): + ranking_filter = gr.Radio(value="All", choices=["All", "1", "2", "3", "4", "5", "None", "aesthetic_score", "random"], label="ranking filter", interactive="true") with gr.Row(): - sort_by = gr.Radio(value="date", choices=["path name", "date", "aesthetic_score", "random"], label="sort by") - sort_order = ToolButton(value=down_symbol) - with gr.Row(): - keyword = gr.Textbox(value="", label="keyword") + aes_filter = gr.Textbox(value="", label="minimum aesthetic_score") with gr.Row(): with gr.Column(): img_file_info = gr.Textbox(label="Generate Info", interactive=False, lines=6) @@ -380,7 +600,8 @@ def create_tab(tabname): except: pass with gr.Row(): - collected_warning = gr.HTML() + collected_warning = gr.HTML() + # hidden items with gr.Row(visible=False): @@ -420,7 +641,10 @@ def create_tab(tabname): end_page.click(lambda s: (-1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) load_switch.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) + exif_keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) + aes_filter.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) sort_by.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) + ranking_filter.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) page_index.submit(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch]) renew_page.click(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch]) refresh_index_button.click(lambda p, s:(p, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch]) @@ -428,7 +652,7 @@ def create_tab(tabname): turn_page_switch.change( fn=get_image_page, - inputs=[img_path, page_index, filenames, keyword, sort_by, sort_order, tabname_box, img_path_depth], + inputs=[img_path, page_index, filenames, keyword, sort_by, sort_order, tabname_box, img_path_depth, ranking_filter, aes_filter, exif_keyword], outputs=[filenames, page_index, history_gallery, img_file_name, img_file_time, img_file_info, visible_img_num, warning_box] ) turn_page_switch.change(fn=None, inputs=[tabname_box], outputs=None, _js="images_history_turnpage") @@ -464,15 +688,49 @@ def create_tab(tabname): # other functions set_index.click(show_image_info, _js="images_history_get_current_img", inputs=[tabname_box, image_index, page_index, filenames], outputs=[img_file_name, img_file_time, image_index, hidden]) - set_index.click(fn=lambda:(gr.update(visible=True), gr.update(visible=True)), inputs=None, outputs=[delete_panel, button_panel]) - img_file_name.change(fn=lambda : "", inputs=None, outputs=[collected_warning]) - - hidden.change(fn=modules.extras.run_pnginfo, inputs=[hidden], outputs=[info1, img_file_info, info2]) + set_index.click(fn=lambda:(gr.update(visible=True), gr.update(visible=True)), inputs=None, outputs=[delete_panel, button_panel]) + img_file_name.change(fn=lambda : "", inputs=None, outputs=[collected_warning]) + img_file_name.change(get_ranking, inputs=img_file_name, outputs=ranking) + + hidden.change(fn=run_pnginfo, inputs=[hidden, img_path, img_file_name], outputs=[info1, img_file_info, info2]) + + #ranking + ranking.change(create_ranked_file, inputs=[img_file_name, ranking]) + #ranking.change(show_next_image_info, _js="images_history_get_current_img", inputs=[tabname_box, image_index, page_index, auto_next], outputs=[img_file_name, img_file_time, image_index, hidden]) + + try: modules.generation_parameters_copypaste.bind_buttons(send_to_buttons, hidden, img_file_info) except: pass + +def run_pnginfo(image, image_path, image_file_name): + if image is None: + return '', '', '' + geninfo, items = images.read_info_from_image(image) + items = {**{'parameters': geninfo}, **items} + + info = '' + for key, text in items.items(): + info += f""" +
+

{plaintext_to_html(str(key))}

+

{plaintext_to_html(str(text))}

+
+""".strip()+"\n" + + if geninfo is None: + try: + filename = os.path.splitext(image_file_name)[0] + ".txt" + geninfo = "" + with open(filename) as f: + for line in f: + geninfo += line + except Exception: + print(f"No EXIF in PNG or txt file") + return '', geninfo, info + def on_ui_tabs(): global num_of_imgs_per_page @@ -480,7 +738,7 @@ def on_ui_tabs(): num_of_imgs_per_page = int(opts.images_history_page_columns * opts.images_history_page_rows) loads_files_num = int(opts.images_history_pages_perload * num_of_imgs_per_page) with gr.Blocks(analytics_enabled=False) as images_history: - with gr.Tabs(elem_id="images_history_tab") as tabs: + with gr.Tabs(elem_id="images_history_tab)") as tabs: for tab in tabs_list: with gr.Tab(tab): with gr.Blocks(analytics_enabled=False) :