a batch label generator for AI Images, useful for generating training data

batch_autolabel.py

import argparse
import glob
import sys
from PIL import Image
import os

def main():
    """
    When given a directory, this script will create text files containing the positive prompts of each image (if the data exists in the png metadata),  
    and stores it as a .txt file with the same name next to the original image.

    This is useful for creating training data.

    Usage:
        python batch_autolabel.py "C:\Users\you\Desktop\training_images"
    """
    #Accept a filepath as an arguments from argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("filepath")
    args = parser.parse_args()
    filepath = args.filepath

    #validate the path
    if not os.path.exists(filepath):
        print("Filepath does not exist")
        sys.exit(1)

    #get all png files that exist in the path
    png_files = glob.glob(filepath + "/*.png")

    #read the png file's metadata
    for png_file in png_files:
        img = Image.open(png_file)
        img.load()  
        #make sure img.info has a 'parameters' key
        if 'parameters' not in img.info:
            continue
        #take the value of the parameters key and store the value up until the '\nNegative' substring, and store it in the variable 'prompt'
        prompt = img.info['parameters'].split('\nNegative')[0]

        #store the 'prompt' string in a file with the same name as the png image, but with the .txt extension
        with open(png_file.split('.')[0] + '.txt', 'w') as f:
            f.write(prompt)

    print("Finished labelling all files.")
        
if __name__ == "__main__":
    main()