Spaces:
Build error
Build error
| # 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() |