init dump with gpt 4o

caption.py

@@ -0,0 +1,102 @@
+import base64
+import io
+import os
+from openai import OpenAI
+from PIL import Image
+from dotenv import load_dotenv
+
+load_dotenv()
+
+def get_prompt():
+    return """Automated Image Captioning (for LoRA Training)
+
+Role: You are an expert AI captioning system generating precise, structured descriptions for AI-generated character images optimized for LoRA model training in Stable Diffusion and Flux.1-dev.
+
+General Guidelines:
+1. Prioritize Consistency – Maintain uniform descriptions across all images in a dataset. Avoid introducing variation in features that should remain constant (e.g., fixed traits like eye color, hair color, or markings that are inherently part of the concept and handled during model training).
+2. Concise and Structured – Only describe visible and significant visual attributes. Use a standardized format for clarity and efficiency.
+3. Omit Subjective Language – Do not evaluative or emotional descriptors like "beautiful" or "scary."
+4. Focus on Key Visual Cues – Clearly describe clothing, accessories, pose, facial expression, lighting, and camera angle. Mention distinctive features only if variable (e.g., visible scar, glasses).
+5. Adapt to Visual Style – Clearly state the artistic style: "anime-style," "photorealistic," "3D-rendered," etc.
+6. Standard Format – Ensure captions follow a consistent structure.
+7. Remain Objective – Do not reference known characters, franchises, or people, even if recognizable. Describe only what is visually present.
+
+Avoid Describing These Unless Variable Across Dataset or Uncertain from Concept:
+- Eye color
+- Hair color
+- Skin tone
+- Tattoos or markings if core to the concept
+- Known accessories that always appear (unless outfit-specific)
+
+Updated Caption Format:  
+tr1gger [Style], [Notable Visual Features], [Clothing], [Pose], [Expression], [Lighting], [Camera Angle]
+
+Captioning Principles:
+- Emphasize visual variation and context-specific details (outfit, pose, lighting, expression, camera angle).
+- Exclude constant traits unless clearly deviating.
+- Include visible clothing and accessories.
+- Clearly define pose and facial expression (neutral, smiling, aggressive).
+- Specify lighting conditions (soft lighting, harsh shadows, glowing backlight).
+- Explicitly state camera angle (e.g., front view, right side profile, low-angle, high-angle, overhead).
+- Avoid mentioning real or fictional identities.
+- Always prefix with the trigger word "tr1gger."
+
+Updated Examples:
+- tr1gger photorealistic, tactical vest and gloves, standing in profile, neutral expression, overcast lighting, side profile
+- tr1gger 3D-rendered, hooded cloak with digital pattern, seated cross-legged, calm expression, low ambient lighting, front view
+- tr1gger anime-style, school uniform with blue necktie, standing with arms behind back, gentle smile, soft daylight, three-quarter view
+- tr1gger photorealistic, long trench coat and combat boots, walking through rain-soaked street, determined expression, dramatic shadows, low-angle view
+"""
+
+def caption_images(images):
+    # Convert PIL images to base64 encoded strings
+    image_strings = []
+    for image in images:
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        image_strings.append(img_str)
+    
+    client = OpenAI()
+    captions = []
+
+    # Start a separate chat session for each image
+    for img_str in image_strings:
+        messages = [
+            {"role": "system", "content": get_prompt()},
+            {
+                "role": "user", 
+                "content": [
+                    {"type": "text", "text": "Caption this image according to the guidelines."},
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}}
+                ]
+            }
+        ]
+        
+        # Request caption for the image in a single chat
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=messages,
+            max_tokens=512  # Adjust max_tokens as needed
+        )
+        
+        # Extract caption from the response
+        caption = response.choices[0].message.content.strip()
+        captions.append(caption)
+    
+    return captions
+
+# Example usage
+if __name__ == "__main__":
+    if not os.environ.get("OPENAI_API_KEY"):
+        print("Please update the .env file with your OpenAI API key.")
+        exit(1)
+
+    # Load images
+    image_paths = ['input/daenyrs_hd.jpg', 'input/girl_body.png']
+    images = [Image.open(path).convert("RGB") for path in image_paths]
+
+    # Generate captions
+    captions = caption_images(images)
+    for i, caption in enumerate(captions):
+        print(f"Generated Caption for Image {i+1}: {caption}")

main.py

@@ -0,0 +1,138 @@
+import os
+import argparse
+import shutil
+import sys
+from pathlib import Path
+from PIL import Image
+from caption import caption_images
+
+def is_image_file(filename):
+    """Check if a file is an allowed image type."""
+    allowed_extensions = ['.png', '.jpg', '.jpeg', '.webp']
+    return any(filename.lower().endswith(ext) for ext in allowed_extensions)
+
+def is_unsupported_image(filename):
+    """Check if a file is an image but not of an allowed type."""
+    unsupported_extensions = ['.bmp', '.gif', '.tiff', '.tif', '.ico', '.svg']
+    return any(filename.lower().endswith(ext) for ext in unsupported_extensions)
+
+def is_text_file(filename):
+    """Check if a file is a text file."""
+    return filename.lower().endswith('.txt')
+
+def validate_input_directory(input_dir):
+    """Validate that the input directory only contains allowed image formats."""
+    input_path = Path(input_dir)
+    
+    unsupported_files = []
+    text_files = []
+    
+    for file_path in input_path.iterdir():
+        if file_path.is_file():
+            if is_unsupported_image(file_path.name):
+                unsupported_files.append(file_path.name)
+            elif is_text_file(file_path.name):
+                text_files.append(file_path.name)
+    
+    if unsupported_files:
+        print("Error: Unsupported image formats detected.")
+        print("Only .png, .jpg, .jpeg, and .webp files are allowed.")
+        print("The following files are not supported:")
+        for file in unsupported_files:
+            print(f"  - {file}")
+        sys.exit(1)
+    
+    if text_files:
+        print("Error: Text files detected in the input directory.")
+        print("The input directory should only contain image files to prevent overwriting existing text files.")
+        print("The following text files were found:")
+        for file in text_files:
+            print(f"  - {file}")
+        sys.exit(1)
+
+def process_images(input_dir, output_dir, fix_outfit=False):
+    """Process all images in the input directory and generate captions."""
+    input_path = Path(input_dir)
+    output_path = Path(output_dir) if output_dir else input_path
+    
+    # Validate the input directory first
+    validate_input_directory(input_dir)
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+    
+    # Track the number of processed images
+    processed_count = 0
+    
+    # Collect all images into a list
+    images = []
+    image_paths = []
+
+    # Get all files in the input directory
+    for file_path in input_path.iterdir():
+        if file_path.is_file() and is_image_file(file_path.name):
+            try:
+                # Load the image
+                image = Image.open(file_path).convert("RGB")
+                images.append(image)
+                image_paths.append(file_path)
+            except Exception as e:
+                print(f"Error loading {file_path.name}: {e}")
+
+    # Log the number of images found
+    print(f"Found {len(images)} images to process.")
+
+    if not images:
+        print("No valid images found to process.")
+        return
+
+    # Generate captions for all images
+    try:
+        captions = caption_images(images)
+    except Exception as e:
+        print(f"Error generating captions: {e}")
+        return
+
+    # Write captions to files
+    for file_path, caption in zip(image_paths, captions):
+        try:
+            # Create caption file path (same name but with .txt extension)
+            caption_filename = file_path.stem + ".txt"
+            caption_path = input_path / caption_filename
+
+            # Write caption to file
+            with open(caption_path, 'w', encoding='utf-8') as f:
+                f.write(caption)
+
+            # If output directory is different from input, copy files
+            if output_path != input_path:
+                # Copy image to output directory
+                shutil.copy2(file_path, output_path / file_path.name)
+                # Copy caption to output directory
+                shutil.copy2(caption_path, output_path / caption_filename)
+
+            processed_count += 1
+            print(f"Processed {file_path.name} → {caption_filename}")
+        except Exception as e:
+            print(f"Error processing {file_path.name}: {e}")
+
+    print(f"\nProcessing complete. {processed_count} images were captioned.")
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate captions for images using GPT-4o.')
+    parser.add_argument('--input', type=str, required=True, help='Directory containing images')
+    parser.add_argument('--output', type=str, help='Directory to save images and captions (defaults to input directory)')
+    parser.add_argument('--fix_outfit', action='store_true', help='Flag to indicate if character has one outfit')
+    
+    args = parser.parse_args()
+    
+    # Validate input directory
+    if not os.path.isdir(args.input):
+        print(f"Error: Input directory '{args.input}' does not exist.")
+        return
+    
+    # Process images
+    process_images(args.input, args.output, args.fix_outfit)
+
+if __name__ == "__main__":
+    main()