import gradio as gr from easygui import msgbox, boolbox from .common_gui import get_folder_path, scriptdir, list_dirs from math import ceil import os import re from .custom_logging import setup_logging # Set up logging log = setup_logging() IMAGES_TO_SHOW = 5 IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") auto_save = True def _get_caption_path(image_file, images_dir, caption_ext): """ Returns the expected path of a caption file for a given image path """ caption_file_name = os.path.splitext(image_file)[0] + caption_ext caption_file_path = os.path.join(images_dir, caption_file_name) return caption_file_path def _get_quick_tags(quick_tags_text): """ Gets a list of tags from the quick tags text box """ quick_tags = [t.strip() for t in quick_tags_text.split(",") if t.strip()] quick_tags_set = set(quick_tags) return quick_tags, quick_tags_set def _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set): """ Updates a list of caption checkboxes to show possible tags and tags already included in the caption """ caption_tags_have = [c.strip() for c in caption.split(",") if c.strip()] caption_tags_unique = [t for t in caption_tags_have if t not in quick_tags_set] caption_tags_all = quick_tags + caption_tags_unique return gr.CheckboxGroup(choices=caption_tags_all, value=caption_tags_have) def paginate_go(page, max_page): try: page = float(page) except: msgbox(f"Invalid page num: {page}") return return paginate(page, max_page, 0) def paginate(page, max_page, page_change): return int(max(min(page + page_change, max_page), 1)) def save_caption(caption, caption_ext, image_file, images_dir): caption_path = _get_caption_path(image_file, images_dir, caption_ext) with open(caption_path, "w+", encoding="utf-8") as f: f.write(caption) log.info(f"Wrote captions to {caption_path}") def update_quick_tags(quick_tags_text, *image_caption_texts): quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text) return [ _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set) for caption in image_caption_texts ] def update_image_caption( quick_tags_text, caption, image_file, images_dir, caption_ext, auto_save ): if auto_save: save_caption(caption, caption_ext, image_file, images_dir) quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text) return _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set) def update_image_tags( quick_tags_text, selected_tags, image_file, images_dir, caption_ext, auto_save, ): # Try to determine order by quick tags quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text) selected_tags_set = set(selected_tags) output_tags = [t for t in quick_tags if t in selected_tags_set] + [ t for t in selected_tags if t not in quick_tags_set ] caption = ", ".join(output_tags) if auto_save: save_caption(caption, caption_ext, image_file, images_dir) return caption def import_tags_from_captions( images_dir, caption_ext, quick_tags_text, ignore_load_tags_word_count ): """ Scans images directory for all available captions and loads all tags under a specified word count into the quick tags box """ def empty_return(): return gr.Text() # Check for images_dir if not images_dir: msgbox("Image folder is missing...") return empty_return() if not os.path.exists(images_dir): msgbox("Image folder does not exist...") return empty_return() if not caption_ext: msgbox("Please provide an extension for the caption files.") return empty_return() if quick_tags_text: if not boolbox( f"Are you sure you wish to overwrite the current quick tags?", choices=("Yes", "No"), ): return empty_return() images_list = os.listdir(images_dir) image_files = [f for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)] # Use a set for lookup but store order with list tags = [] tags_set = set() for image_file in image_files: caption_file_path = _get_caption_path(image_file, images_dir, caption_ext) if os.path.exists(caption_file_path): with open(caption_file_path, "r", encoding="utf-8") as f: caption = f.read() for tag in caption.split(","): tag = tag.strip() tag_key = tag.lower() if not tag_key in tags_set: # Ignore extra spaces total_words = len(re.findall(r"\s+", tag)) + 1 if total_words <= ignore_load_tags_word_count: tags.append(tag) tags_set.add(tag_key) return ", ".join(tags) def load_images(images_dir, caption_ext, loaded_images_dir, page, max_page): """ Triggered to load a new set of images from the folder to caption This loads in the total expected image counts to be used by pagination before running update_images """ def empty_return(): return [loaded_images_dir, page, max_page] # Check for images_dir if not images_dir: msgbox("Image folder is missing...") return empty_return() if not os.path.exists(images_dir): msgbox("Image folder does not exist...") return empty_return() if not caption_ext: msgbox("Please provide an extension for the caption files.") return empty_return() # Load Images images_list = os.listdir(images_dir) total_images = len( [True for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)] ) return [images_dir, 1, ceil(total_images / IMAGES_TO_SHOW)] def update_images( images_dir, caption_ext, quick_tags_text, page, ): """ Updates the displayed images and captions from the current page and image directory """ # Load Images images_list = os.listdir(images_dir) image_files = [f for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)] # Quick tags quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text or "") # Display Images rows = [] image_paths = [] captions = [] tag_checkbox_groups = [] start_index = (int(page) - 1) * IMAGES_TO_SHOW for i in range(IMAGES_TO_SHOW): image_index = start_index + i show_row = image_index < len(image_files) image_path = None caption = "" tag_checkboxes = None if show_row: image_file = image_files[image_index] image_path = os.path.join(images_dir, image_file) caption_file_path = _get_caption_path(image_file, images_dir, caption_ext) if os.path.exists(caption_file_path): with open(caption_file_path, "r", encoding="utf-8") as f: caption = f.read() tag_checkboxes = _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set) rows.append(gr.Row(visible=show_row)) image_paths.append(image_path) captions.append(caption) tag_checkbox_groups.append(tag_checkboxes) return ( rows + image_paths + image_paths + captions + tag_checkbox_groups + [gr.Row(visible=True), gr.Row(visible=True)] ) # Gradio UI def gradio_manual_caption_gui_tab(headless=False, default_images_dir=None): from .common_gui import create_refresh_button default_images_dir = ( default_images_dir if default_images_dir is not None else os.path.join(scriptdir, "data") ) current_images_dir = default_images_dir # Function to list directories def list_images_dirs(path): # Allows list_images_dirs to modify current_images_dir outside of this function nonlocal current_images_dir current_images_dir = path return list(list_dirs(path)) with gr.Tab("Manual Captioning"): gr.Markdown("This utility allows quick captioning and tagging of images.") page = gr.Number(value=-1, visible=False) max_page = gr.Number(value=1, visible=False) loaded_images_dir = gr.Text(visible=False) with gr.Group(), gr.Row(): images_dir = gr.Dropdown( label="Image folder to caption (containing the images to caption)", choices=[""] + list_images_dirs(default_images_dir), value="", interactive=True, allow_custom_value=True, ) create_refresh_button( images_dir, lambda: None, lambda: {"choices": list_images_dirs(current_images_dir)}, "open_folder_small", ) folder_button = gr.Button( "📂", elem_id="open_folder_small", elem_classes=["tool"], visible=(not headless), ) folder_button.click( get_folder_path, outputs=images_dir, show_progress=False, ) load_images_button = gr.Button("Load", elem_id="open_folder") caption_ext = gr.Dropdown( label="Caption file extension", choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, allow_custom_value=True, ) auto_save = gr.Checkbox( label="Autosave", info="Options", value=True, interactive=True ) images_dir.change( fn=lambda path: gr.Dropdown(choices=[""] + list_images_dirs(path)), inputs=images_dir, outputs=images_dir, show_progress=False, ) # Caption Section with gr.Group(), gr.Row(): quick_tags_text = gr.Textbox( label="Quick Tags", placeholder="Comma separated list of tags", interactive=True, ) import_tags_button = gr.Button("Import", elem_id="open_folder") ignore_load_tags_word_count = gr.Slider( minimum=1, maximum=100, value=3, step=1, label="Ignore Imported Tags Above Word Count", interactive=True, ) # Next/Prev section generator def render_pagination(): gr.Button("< Prev", elem_id="open_folder").click( paginate, inputs=[page, max_page, gr.Number(value=-1, visible=False)], outputs=[page], ) page_count = gr.Label("Page 1", label="Page") page_goto_text = gr.Textbox( label="Goto page", placeholder="Page Number", interactive=True, ) gr.Button("Go >", elem_id="open_folder").click( paginate_go, inputs=[page_goto_text, max_page], outputs=[page], ) gr.Button("Next >", elem_id="open_folder").click( paginate, inputs=[page, max_page, gr.Number(value=1, visible=False)], outputs=[page], ) return page_count with gr.Row(visible=False) as pagination_row1: page_count1 = render_pagination() # Images section image_rows = [] image_files = [] image_images = [] image_caption_texts = [] image_tag_checks = [] save_buttons = [] for _ in range(IMAGES_TO_SHOW): with gr.Row(visible=False) as row: image_file = gr.Text(visible=False) image_files.append(image_file) image_image = gr.Image(type="filepath") image_images.append(image_image) image_caption_text = gr.TextArea( label="Captions", placeholder="Input captions", interactive=True, ) image_caption_texts.append(image_caption_text) tag_checkboxes = gr.CheckboxGroup([], label="Tags", interactive=True) save_button = gr.Button( "💾", elem_id="open_folder_small", elem_classes=["tool"], visible=False, ) save_buttons.append(save_button) # Caption text change image_caption_text.input( update_image_caption, inputs=[ quick_tags_text, image_caption_text, image_file, loaded_images_dir, caption_ext, auto_save, ], outputs=tag_checkboxes, ) # Quick tag check tag_checkboxes.input( update_image_tags, inputs=[ quick_tags_text, tag_checkboxes, image_file, loaded_images_dir, caption_ext, auto_save, ], outputs=[image_caption_text], ) # Save Button save_button.click( save_caption, inputs=[ image_caption_text, caption_ext, image_file, images_dir, ], ) image_tag_checks.append(tag_checkboxes) image_rows.append(row) # Next/Prev Section with gr.Row(visible=False) as pagination_row2: page_count2 = render_pagination() # Quick tag text update quick_tags_text.change( update_quick_tags, inputs=[quick_tags_text] + image_caption_texts, outputs=image_tag_checks, ) # Import tags button import_tags_button.click( import_tags_from_captions, inputs=[ loaded_images_dir, caption_ext, quick_tags_text, ignore_load_tags_word_count, ], outputs=quick_tags_text, ) # Load Images button load_images_button.click( load_images, inputs=[ images_dir, caption_ext, loaded_images_dir, page, max_page, ], outputs=[loaded_images_dir, page, max_page], ) # Update images shown when the update key changes # This allows us to trigger a change from multiple # sources (page, image_dir) image_update_key = gr.Text(visible=False) image_update_key.change( update_images, inputs=[loaded_images_dir, caption_ext, quick_tags_text, page], outputs=image_rows + image_files + image_images + image_caption_texts + image_tag_checks + [pagination_row1, pagination_row2], show_progress=False, ) # Update the key on page and image dir change listener_kwargs = { "fn": lambda p, i: f"{p}-{i}", "inputs": [page, loaded_images_dir], "outputs": image_update_key, } page.change(**listener_kwargs) loaded_images_dir.change(**listener_kwargs) # Save buttons visibility # (on auto-save on/off) auto_save.change( lambda auto_save: [gr.Button(visible=not auto_save)] * IMAGES_TO_SHOW, inputs=auto_save, outputs=save_buttons, ) # Page Count page.change( lambda page, max_page: [f"Page {int(page)} / {int(max_page)}"] * 2, inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False, )