File size: 4,785 Bytes
cff1674 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
# 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)
|