Spaces:
Build error
Build error
| # Proposed by https://github.com/kainatquaderee | |
| import os | |
| import argparse | |
| import logging | |
| from pathlib import Path | |
| def is_image_file(filename, image_extensions): | |
| """Check if a file is an image file based on its extension.""" | |
| return Path(filename).suffix.lower() in image_extensions | |
| def create_text_file(image_filename, output_directory, text_extension): | |
| """Create a text file with the same name as the image file.""" | |
| # Extract prompt from filename | |
| prompt = Path(image_filename).stem | |
| # Construct path for the output text file | |
| text_file_path = Path(output_directory) / (prompt + text_extension) | |
| try: | |
| # Write prompt to text file | |
| with open(text_file_path, 'w') as text_file: | |
| text_file.write(prompt) | |
| logging.info(f"Text file created: {text_file_path}") | |
| return 1 | |
| except IOError as e: | |
| logging.error(f"Failed to write to {text_file_path}: {e}") | |
| return 0 | |
| def main(image_directory, output_directory, image_extension, text_extension): | |
| # If no output directory is provided, use the image directory | |
| if not output_directory: | |
| output_directory = image_directory | |
| # Ensure the output directory exists, create it if necessary | |
| Path(output_directory).mkdir(parents=True, exist_ok=True) | |
| # Initialize a counter for the number of text files created | |
| text_files_created = 0 | |
| # Iterate through files in the directory | |
| for image_filename in Path(image_directory).iterdir(): | |
| # Check if the file is an image | |
| if is_image_file(image_filename, image_extension): | |
| # Create a text file with the same name as the image file and increment the counter if successful | |
| text_files_created += create_text_file(image_filename, output_directory, text_extension) | |
| # Report if no text files were created | |
| if text_files_created == 0: | |
| logging.info("No image matching extensions were found in the specified directory. No caption files were created.") | |
| else: | |
| logging.info(f"{text_files_created} text files created successfully.") | |
| def create_gui(image_directory, output_directory, image_extension, text_extension): | |
| try: | |
| import gradio | |
| import gradio.blocks as blocks | |
| except ImportError: | |
| print("gradio module is not installed. Please install it to use the GUI.") | |
| exit(1) | |
| """Create a Gradio interface for the caption creation process.""" | |
| with gradio.Blocks() as demo: | |
| gradio.Markdown("## Caption From Filename") | |
| with gradio.Row(): | |
| with gradio.Column(): | |
| image_dir = gradio.Textbox(label="Image Directory", value=image_directory) | |
| output_dir = gradio.Textbox(label="Output Directory", value=output_directory) | |
| image_ext = gradio.Textbox(label="Image Extensions", value=" ".join(image_extension)) | |
| text_ext = gradio.Textbox(label="Text Extension", value=text_extension) | |
| run_button = gradio.Button("Run") | |
| with gradio.Column(): | |
| output = gradio.Textbox(label="Output", placeholder="Output will be displayed here...", lines=10, max_lines=10) | |
| run_button.click(main, inputs=[image_dir, output_dir, image_ext, text_ext], outputs=output) | |
| demo.launch() | |
| if __name__ == "__main__": | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') | |
| # Create an argument parser | |
| parser = argparse.ArgumentParser(description='Generate caption files from image filenames.') | |
| # Add arguments for the image directory, output directory, and file extension | |
| parser.add_argument('image_directory', help='Directory containing the image files.') | |
| parser.add_argument('--output_directory', help='Optional: Output directory where text files will be saved. If not provided, the files will be saved in the same directory as the images.') | |
| parser.add_argument('--image_extension', nargs='+', default=['.jpg', '.jpeg', '.png', '.webp', '.bmp'], help='Extension(s) for the image files. Defaults to common image extensions .jpg, .jpeg, .png, .webp, .bmp.') | |
| parser.add_argument('--text_extension', default='.txt', help='Extension for the output text files. Defaults to .txt.') | |
| parser.add_argument('--gui', action='store_true', help='Launch a Gradio interface for the caption creation process.') | |
| # Parse the command-line arguments | |
| args = parser.parse_args() | |
| if args.gui: | |
| create_gui(args.image_directory, args.output_directory, args.image_extension, args.text_extension) | |
| else: | |
| main(args.image_directory, args.output_directory, args.image_extension, args.text_extension) | |