concat_captions and zsh fuckery

Signed-off-by: Balazs Horvath <acsipont@gmail.com>

.zshrc

@@ -261,17 +261,17 @@ function re() {
 # It searches for all .txt files in the specified directory and its subdirectories.
 # If a file contains the specified tag, the function removes the tag from its original position
 # and prepends it to the beginning of the file.
-# Usage: rejiggle_captions <tag> <directory>
-rejiggle_captions() {
+# Usage: rejiggle_tags <tag> <directory>
+rejiggle_tags() {
   local tag="$1"
   local dir="$2"
 
   if [[ -z "$tag" || -z "$dir" ]]; then
-    echo "Usage: rejiggle_captions <tag> <directory>"
+    echo "Usage: rejiggle_tags <tag> <directory>"
     return 1
   fi
 
-  find "$dir" -type f -name "*.txt" | while read -r file; do
+  find "$dir" -type f -name "*.tags" | while read -r file; do
     if grep -q "$tag" "$file"; then
       sed -i "s/$tag//g" "$file"
       sed -i "1s/^/$tag, /" "$file"
@@ -385,14 +385,14 @@ list_word_freqs() {
 
 # Function: sample_prompts
 # Description:
-#   This function takes a sample of the tags (or captions) in a target training directory.
+#   This function takes a sample of the tag in a target training directory.
 #   It reads and displays the contents of all .txt files in the specified directory,
-#   providing a quick overview of the tags or captions used for training.
+#   providing a quick overview of the tags used for training.
 #
 # Usage: sample_prompts <target_directory>
 #
 # Parameters:
-#   - target_directory: The directory containing the .txt files with tags or captions.
+#   - target_directory: The directory containing the .txt files with tags.
 #
 # Output:
 #   Prints the contents of each .txt file in the target directory, separated by newlines.
@@ -406,12 +406,12 @@ sample_prompts() {
 }
 
 # replace_comma_with_keep_tags
-# Description: This function replaces the specified occurrence of a comma with " |||" in all *.txt files
+# Description: This function replaces the specified occurrence of a comma with " |||" in all *.tags files
 #              in all subdirectories of a target directory or the current directory when no path is passed.
 # Usage: replace_comma_with_keep_tags <occurrence_number> [target_directory]
 # Parameters:
 #   - occurrence_number: The occurrence number of the comma to be replaced (e.g., 1 for the first occurrence).
-#   - target_directory (optional): The target directory to search for *.txt files. If not provided, the current directory is used.
+#   - target_directory (optional): The target directory to search for *.tags files. If not provided, the current directory is used.
 # Example:
 #   replace_comma_with_keep_tags 2 /path/to/directory
 #   replace_comma_with_keep_tags 1
@@ -424,7 +424,7 @@ replace_comma_with_keep_tags() {
     return 1
   fi
 
-  find "$target_directory" -type f -name "*.txt" | while read -r file; do
+  find "$target_directory" -type f -name "*.tags" | while read -r file; do
     awk -v occurrence="$occurrence_number" '{
       count = 0
       for (i = 1; i <= NF; i++) {
@@ -670,18 +670,18 @@ replace_text_in_files() {
 
 # This script adds a specified prefix to the beginning of each text file in a given directory.
 # If the prefix already exists in the text file, it moves the prefix to the front of the text file without leaving extra commas or spaces.
-# Usage: inject_to_captions <directory> <prefix>
+# Usage: inject_to_tags <directory> <prefix>
 # Arguments:
 #   <directory> - The directory containing the text files to be modified.
 #   <prefix> - The prefix to be added to the beginning of each text file.
 # The script checks if the specified directory exists and iterates over each text file in the directory.
 # For each text file, it creates a temporary file with the modified content and then replaces the original file with the temporary file.
 # If the directory does not exist, it prints an error message.
-inject_to_captions() {
+inject_to_tags() {
     local dir="$1"
     local prefix="$2"
     if [[ -d "$dir" ]]; then
-        for file in "$dir"/*.txt; do
+        for file in "$dir"/*.tags; do
             if [[ -f "$file" ]]; then
                 if grep -q "$prefix" "$file"; then
                     # Move the existing prefix to the front of the text file without leaving extra commas or spaces

concat_captions

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+"""
+This script walks through a directory, identifies image files, and checks for the existence of corresponding
+.caption and .tags files. It then concatenates the contents of .caption and .tags files into the .txt files.
+
+Usage:
+    - Place the script in the directory containing the image files.
+    - Run the script to concatenate .caption and .tags files into .txt files.
+    - Use the dry_run flag to preview the changes without writing to the .txt files.
+
+Functions:
+    get_files(path): Walks through the directory and yields image files along with their .caption and .tags files.
+    concat(caption_path, tags_path, txt_path, dry_run=False): Concatenates the contents of .caption and .tags files into the .txt file.
+"""
+
+from pathlib import Path
+import os
+
+FILE_EXTS = {".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".jxl"}
+
+
+def get_files(path):
+    path = Path(path)
+    # Walk the directory, looking for image files
+    for root, dirs, files in os.walk(path):
+        root = path / root
+        for file in files:
+            file = root / file
+            if file.suffix not in FILE_EXTS:
+                continue
+            caption = file.with_suffix(".caption")
+            tags = file.with_suffix(".tags")
+            txt = file.with_suffix(".txt")
+            if not caption.exists():
+                print(f"{caption} does not exist")
+            if not tags.exists():
+                print(f"{tags} does not exist")
+            yield file, caption, tags, txt
+
+
+def concat(caption_path, tags_path, txt_path, dry_run=False):
+    with open(caption_path, "r") as f:
+        caption = f.read().strip()
+
+    with open(tags_path, "r") as f:
+        tags = f.read().strip(", \n")
+
+    txt = f"{tags}, {caption}"
+
+    if dry_run:
+        print(f"{txt_path}:")
+        print(txt)
+        print()
+    else:
+        with open(txt_path, 'w') as f:
+            f.write(txt)
+        print(f"wrote {txt_path}")
+
+if __name__ == "__main__":
+    dry_run = False
+    for f in get_files("."):
+        concat(*f[1:], dry_run=dry_run)
+