File size: 3,127 Bytes
fe6327d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
# This script will create the caption text files in the specified folder using the specified file pattern and caption text.
#
# eg: python caption.py D:\some\folder\location "*.png, *.jpg, *.webp" "some caption text"
import argparse
# import glob
# import os
from pathlib import Path
def create_caption_files(image_folder: str, file_pattern: str, caption_text: str, caption_file_ext: str, overwrite: bool):
# Split the file patterns string and strip whitespace from each pattern
patterns = [pattern.strip() for pattern in file_pattern.split(",")]
# Create a Path object for the image folder
folder = Path(image_folder)
# Iterate over the file patterns
for pattern in patterns:
# Use the glob method to match the file patterns
files = folder.glob(pattern)
# Iterate over the matched files
for file in files:
# Check if a text file with the same name as the current file exists in the folder
txt_file = file.with_suffix(caption_file_ext)
if not txt_file.exists() or overwrite:
# Create a text file with the caption text in the folder, if it does not already exist
# or if the overwrite argument is True
with open(txt_file, "w") as f:
f.write(caption_text)
def main():
# Define command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument("image_folder", type=str, help="the folder where the image files are located")
parser.add_argument("--file_pattern", type=str, default="*.png, *.jpg, *.jpeg, *.webp", help="the pattern to match the image file names")
parser.add_argument("--caption_file_ext", type=str, default=".caption", help="the caption file extension.")
parser.add_argument("--overwrite", action="store_true", default=False, help="whether to overwrite existing caption files")
# Create a mutually exclusive group for the caption_text and caption_file arguments
group = parser.add_mutually_exclusive_group()
group.add_argument("--caption_text", type=str, help="the text to include in the caption files")
group.add_argument("--caption_file", type=argparse.FileType("r"), help="the file containing the text to include in the caption files")
# Parse the command-line arguments
args = parser.parse_args()
image_folder = args.image_folder
file_pattern = args.file_pattern
caption_file_ext = args.caption_file_ext
overwrite = args.overwrite
# Get the caption text from either the caption_text or caption_file argument
if args.caption_text:
caption_text = args.caption_text
elif args.caption_file:
caption_text = args.caption_file.read()
# Create a Path object for the image folder
folder = Path(image_folder)
# Check if the image folder exists and is a directory
if not folder.is_dir():
raise ValueError(f"{image_folder} is not a valid directory.")
# Create the caption files
create_caption_files(image_folder, file_pattern, caption_text, caption_file_ext, overwrite)
if __name__ == "__main__":
main() |