69 lines
3.1 KiB
Python
69 lines
3.1 KiB
Python
# 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() |