utils/concat_captions.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# 画像ファイルと.captionと.tagsファイルを.txtファイルに連結する
# Concatenate image files and .caption and .tags files to .txt files
#
# This script walks through a directory, identifies image files, and checks for the existence of corresponding
# .caption and .tags files. It then concatenates the contents of .caption and .tags files into the .txt files.
#
# Usage:
#     - Place the script in the directory containing the image files.
#     - Run the script to concatenate .caption and .tags files into .txt files.
#     - Use the dry_run flag to preview the changes without writing to the .txt files.
#     - Files missing either .caption or .tags will be skipped.
#
# Functions:
#     get_files(path): Walks through the directory and yields image files along with their .caption and .tags files.
#     concat(caption_path, tags_path, txt_path, dry_run=False): Concatenates the contents of .caption and .tags files into the .txt file.

from pathlib import Path
import os

FILE_EXTS = {".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".jxl", ".webp"}


def get_files(path):
    path = Path(path)
    # Walk the directory, looking for image files
    for root, dirs, files in os.walk(path):
        root = path / root
        for file in files:
            file = root / file
            if file.suffix not in FILE_EXTS:
                continue
            caption = file.with_suffix(".caption")
            tags = file.with_suffix(".tags")
            txt = file.with_suffix(".txt")
            
            if not caption.exists():
                print(f"Skipping {file}: {caption} does not exist")
                continue
            if not tags.exists():
                print(f"Skipping {file}: {tags} does not exist")
                continue
                
            yield file, caption, tags, txt


def concat(caption_path, tags_path, txt_path, dry_run=False):
    with open(caption_path, "r") as f:
        caption = f.read().strip()

    with open(tags_path, "r") as f:
        tags = f.read().strip(", \n")

    txt = f"{tags}, {caption}"

    if dry_run:
        print(f"{txt_path}:")
        print(txt)
        print()
    else:
        with open(txt_path, 'w') as f:
            f.write(txt)
        print(f"wrote {txt_path}")

if __name__ == "__main__":
    dry_run = False
    for f in get_files("."):
        concat(*f[1:], dry_run=dry_run)