🚶

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

に公開
AI
Python
pillow
生成 AI
ComfyUI
tech

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を含めたキーワード検索を行いますので、キーワードはなるべく正確なものを入力してください。

Discussion