From 9053bcd6d8979fd060e69caa2bfaba43ab1b0769 Mon Sep 17 00:00:00 2001 From: Anthony Wu Date: Wed, 22 Jan 2025 13:25:59 -0800 Subject: [PATCH] Batch Image Renamer tool as an isolated uv run script --- README.md | 19 ++++ src/mflux/post_processing/image_util.py | 2 +- tools/rename_images.py | 131 ++++++++++++++++++++++++ 3 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 tools/rename_images.py diff --git a/README.md b/README.md index 9538358..be18791 100644 --- a/README.md +++ b/README.md @@ -777,6 +777,25 @@ with different prompts and LoRA adapters active. - Currently, the supported controlnet is the [canny-only version](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny). - Dreambooth training currently does not support sending in training parameters as flags. +### Optional Tool: Batch Image Renamer + +With a large number of generated images, some users want to automatically rename their image outputs to reflect the prompts and configs. + +The bundled `tools/rename_images.py` is an optional tool that is part with the project repo +but *not* included in the `mflux` Python package due to additional dependencies that +do not make sense to become standard project requirements. + +The script uses [KeyBERT](https://maartengr.github.io/KeyBERT/) (a keyword extraction library) +to extract keywords from mflux exif metadata to update the image file names. +We then use `uv run` to execute the script in an isolated env without affecting your `mflux` env. + +Users who want to use or extend this tool to their own needs is encouraged to `git clone` the repo +then `uv run tools/rename_images.py ` or download the single-file standalone script +and `uv run your/path/rename_images.py`. + +This script's renaming logic can be customized to your needs. +See `uv run tools/rename_images.py --help` for full CLI usage help. + ### 💡Workflow Tips - To hide the model fetching status progress bars, `export HF_HUB_DISABLE_PROGRESS_BARS=1` diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index c232cc4..c7ccca4 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -154,7 +154,7 @@ class ImageUtil: def _embed_metadata(metadata: dict, path: str) -> None: try: # Convert metadata dictionary to a string - metadata_str = str(metadata) + metadata_str = json.dumps(metadata) # Convert the string to bytes (using UTF-8 encoding) user_comment_bytes = metadata_str.encode("utf-8") diff --git a/tools/rename_images.py b/tools/rename_images.py new file mode 100644 index 0000000..89e737b --- /dev/null +++ b/tools/rename_images.py @@ -0,0 +1,131 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "Pillow", +# "keybert", +# ] +# /// + +import ast +import json +from pathlib import Path + +from keybert import KeyBERT +from PIL import Image, UnidentifiedImageError + +PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +EXIF_USER_COMMENT_KEY = 37510 + + +class UnsupportedMetadata(Exception): + pass + + +def parse_metadata_from_image(image: Image, insecure=False) -> dict: + try: + exif_data = image._getexif() + except AttributeError: + # when exif data is not available for some image type/format + return None + + metadata = None + if exif_data: + try: + for tag, value in exif_data.items(): + if tag == EXIF_USER_COMMENT_KEY: + metadata = json.loads(value.decode()) + except KeyError: + metadata = None + except json.decoder.JSONDecodeError: + if not insecure: + raise UnsupportedMetadata( + "The metadata is likely stored as a str literal of a Python dict before mflux 0.6.0. " + "This tool cannot guarantee safety of running eval(...) on the value. " + "However, you can bypass this caution by passing the --insecure flag. " + f"Metadata: {value.decode()}" + ) + try: + # try to parse the str(dict) data + obj = ast.literal_eval(value.decode()) + if isinstance(obj, dict): + metadata = obj + else: + raise UnsupportedMetadata( + "The metadata is not a dict output recognized by this tool. " f"Metadata: {value.decode()}" + ) + except (ValueError, SyntaxError): + raise UnsupportedMetadata("The metadata is not parseable by this tool. " f"Metadata: {value.decode()}") + + return metadata + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="Use mflux metadata to rename images.") + parser.add_argument("paths", nargs="+", help="mflux image files or directories to process") + parser.add_argument( + "--n-keywords", type=int, default=5, help="N number of keywords to extract from each prompt. Default 5." + ) + parser.add_argument( + "--yes", action="store_true", default=False, help="Allow renaming without interactive confirmation." + ) + parser.add_argument( + "--insecure", + action="store_true", + default=False, + help="At your own risk, allow insecure parsing of literal Python values saved by mflux versions < 0.6.0", + ) + args = parser.parse_args() + + if not args.yes: + print( + "INFO: This tool by default dry runs the file rename. To automatically accept renames, pass the `--yes` flag." + ) + + processed_paths = [] + for path in args.paths: + p = Path(path) + if p.is_file(): + processed_paths.append(p) + elif p.is_dir(): + processed_paths.extend(p.glob("*")) + + keyword_model = KeyBERT() + unsupported_errors: dict[str, UnsupportedMetadata] = {} + for p in processed_paths: + if p.is_dir(): + continue + try: + with Image.open(p) as image: + image.verify() # Verify that it is an image + except UnidentifiedImageError: + # expected when receiving dir paths, ignore all non-images in dirs + continue + + try: + mflux_metadata = parse_metadata_from_image(image, insecure=args.insecure) + if not mflux_metadata: + unsupported_errors[p.as_posix()] = "metadata not stored in EXIF" + continue + prompt_keywords = keyword_model.extract_keywords(mflux_metadata["prompt"], top_n=args.n_keywords) + proposed_new_stem = f"{'_'.join([word for word, _ in prompt_keywords])}__seed_{mflux_metadata['seed']}" + new_path = p.with_stem(proposed_new_stem) + if args.yes: + old_path = p + p.rename(new_path) + print(f"File renamed: {old_path} -> {new_path}") + else: + print(f"File rename proposed: {p} -> {new_path}") + except UnsupportedMetadata as ume: + unsupported_errors[p.as_posix()] = ume + + for path, error in unsupported_errors.items(): + print(f"{path}: {error}") + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + pass