create-captions-for-directory.py

import os
import sys
import json
from PIL import Image

def main():
  # read command line parameter
  if len(sys.argv) < 2:
    print("Usage: python create-captions-for-directory.py <image_path>")
    return
  imagePath = sys.argv[1]
  prompt = get_prompt_from_image(imagePath)
  print(prompt)

def get_prompt_from_image(imagePath):
  # load the image
  image = Image.open(imagePath)
  # summarize some details about the image
  info = (image.info or {}).copy()
  # check if info has prompt
  if 'prompt' in info:
    # ComfyUI
    prompt = info['prompt']
    # try to parse this json string
    try:
      info = json.loads(prompt)
      # loop through the properties of the json object
      # just keep the inputs with type 'text'
      textNodes = []
      for key in info:
        node = info[key]
        # check if this node has an inputs property
        hasInputsProperty = 'inputs' in node
        hasText = hasInputsProperty and 'text' in node['inputs']
        hasClassTypeProperty = 'class_type' in node
        isTextNode = hasText and hasClassTypeProperty
        if isTextNode:
          textNodes.append(node)
        # print(node['inputs'])
    except:
      return ""
    result = ""
    for node in textNodes:
      result += node['inputs']['text'] + '\n'
    # get the first line
    first_line = result.split('\n')[0]
    return first_line
  # webui
  parameters = info.pop('parameters', None)
  # get the first line of the paramters
  first_line = parameters.split('\n')[0]
  return first_line

if __name__ == "__main__":
  main()