ComfyUIで生成したPNG画像のメタデータを確認してファイル移動する

ComfyUIのような画像生成AIを使っていると、ものすごい数の画像が生成され、PNGファイルが数百枚数千枚と増えていき、分類に困ってしまいます。

ComfyUIで生成したPNG画像にはpromptやworkflowといったJSONメタデータが含まれているので、入力したキーワードごとにフォルダを作成し、キーワードがメタデータに含まれている場合ファイル移動してくれるツールを作りました。

ファイル移動は対象ディレクトリ内に存在する同一名のファイルを上書きしたりするような破壊的行為ですので、スクリプトを実行する前に入力したコマンドに問題がないかしっかり確認しましょう。

import json
import os
import shutil
import sys

from PIL import Image


def contains_keyword(info, keyword):
    """Check if the keyword exists in 'prompt' or 'workflow' fields of image metadata (case-insensitive)"""
    for key in ("prompt", "workflow"):
        if key not in info:
            continue

        try:
            data = json.loads(info[key])
        except Exception:
            continue

        if search_keyword_in_structure(data, keyword):
            return True

    return False


def search_keyword_in_structure(data, keyword):
    """Recursively search for the keyword in a nested structure"""
    if isinstance(data, dict):
        for k, v in data.items():
            if search_keyword_in_structure(k, keyword) or search_keyword_in_structure(v, keyword):
                return True
    elif isinstance(data, list):
        for item in data:
            if search_keyword_in_structure(item, keyword):
                return True
    elif isinstance(data, str):
        if keyword in data.lower():
            return True

    return False


def process_images(target_dir, keyword):
    """Move PNG images containing the keyword in their metadata to a subdirectory named after the keyword."""
    dest_dir = os.path.join(target_dir, keyword)
    os.makedirs(dest_dir, exist_ok=True)

    for filename in os.listdir(target_dir):
        if not filename.lower().endswith(".png"):
            continue

        file_path = os.path.join(target_dir, filename)
        if not os.path.isfile(file_path):
            continue

        try:
            with Image.open(file_path) as img:
                has_keyword = contains_keyword(img.info, keyword.lower())

            if has_keyword:
                shutil.move(file_path, os.path.join(dest_dir, filename))
                print(f"Moved: {filename}")
        except Exception as e:
            print(f"Skipped: {filename} (Error: {e})")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python script.py <target_dir> <keyword>")
        sys.exit(1)

    target_dir = sys.argv[1]
    keyword = sys.argv[2]

    if not os.path.isdir(target_dir):
        print(f"Error: '{target_dir}' is not a valid directory.")
        sys.exit(1)

    process_images(target_dir, keyword)

PythonとPILを使うので、もしComfyUIを使用している環境や、画像を管理している環境にPILがない場合はpipでインストールします。

$ pip install pillow

その後、今回のスクリプトを実行します。

$ python script.py images maid

こうすることでimagesディレクトリにある画像のうち、プロンプトなどで"maid"を使用した画像がimages/maidディレクトリに移動します。

カレントディレクトリにある"school uniform"をプロンプトなどで使用した画像を./school uniformディレクトリに移動したい場合は以下のように実行します。

$ python script.py . "school uniform"

JSONメタデータのValueだけでなく、Keyを含めたキーワード検索を行いますので、キーワードはなるべく正確なものを入力してください。