import os
import base64
import requests
import json

# Folder path
FOLDER_PATH = r'D:\Pictures\data'

# API details
API_URL = 'https://api.x.ai/v1/chat/completions'
API_KEY = 'YOUR_KEY'

# Supported image extensions
IMAGE_EXTENSIONS = ['.png', '.jpeg', '.jpg', '.gif', '.webp']

def encode_image(image_path):
    """Encode image to base64."""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def generate_caption(image_path):
    """Generate caption for an image using X.AI API."""
    base64_image = encode_image(image_path)

    payload = {
        "messages": [
            {
                "role": "system",
                "content": "You are a Danbooru AI bot for auto-captioning. Make precise captions for the provided image, describing its content in format of Danbooru/Safebooru tags only (format example: \"1girl, solo, brown hair, nude, white background, simple background\"), with minimum of 5 and maximum of 99 tags depending on image content. **Follow the rules.** Only use tags that are applicable for the image, skipping the rest. You must always provide Booru tags for the background, like \"indoors\" or \"outdoors\" or \"simple background, white background\" or \"gradient background, grey background\" or \"black background\" etc. For detailed scenery, add tag \"scenery\". For items placed like in a reference sheet, add \"reference sheet\". For a close-up image of a single item, use \"game item icon\", and if there are multiple objects then use \"game assets\". You must also describe presence of persons in the image. If no characters, then \"no humans\"; otherwise use format like \"1girl\"/\"2girls\"/\"multiple girls\"/\"1girl, 1boy\"/\"1other\" (if gender is unclear), etc. If a monster present on image then use tag \"monster\", if an animal then use \"animal\", and for non-monstrous non-human creatures use \"creature\". Describe species whenever possible or use tags like \"horns\", \"pointy ears\", \"antlers\", \"wings\", \"fangs\", \"claws\", \"bat wings\", \"demon wings\" etc. Always specify hair length (\"short hair\", \"medium hair\", \"long hair\", \"very long hair\", \"bald\" etc.), for female characters also specify their breasts size (\"flat chest\", \"small breasts\", \"medium breasts\", \"large breasts\" etc. If some color strongly prevails, add a tag for theme, like \"blue theme\", \"white theme\" etc, but normally don't include it. For a bleak (barely colored) image use \"muted colors\" tag. For fully-monochrome, add \"monochrome\" and for grayscale add \"grayscale\". Those tags are applicable for possible different techniques or artstyles if you found signs of any: smooth shading, digital painting, textured stroke, glossy surfaces, lineart, sketch, outline, flat color, cinematic, etc. Describe the pose by action for legs (if visible) like \"standing\"/\"sitting\"/\"squatting\" etc, and for arms/hands (examples: \"arm up\", \"arms up\", \"arm behind head\", \"croseed arms\", \"hand on hip\"). Describe the direction of gaze for a character (\"looking at viewer\", \"looking away\", \"looking back\", \"looking up\", \"looking to side\" etc). When applicable, describe the angle (\"from behind\", \"from side\", \"from above\" etc). Generally, anything unusual or non-trivial shall have a suitable Booru tag applied. If some of elements seems deeply interconnected, you may use \"and\"/\"with\" between them instead of a comma (like in \"pendant with red gem\", \"1girl and unicorn\"), but don't use natural language or entire sentences. If a character is holding/wielding some item, write a tag literally like \"holding sword\", \"holding chalice\", \"holding phone\" etc. If a text present on image, only use tags like \"signature\", \"copyright name\", \"logo\" etc without writing the specific text content. **Write now only the required line of tags for provided image, and nothing else.**"
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    }
                ]
            }
        ],
        "model": "grok-vision-beta",
        "stream": False,
        "temperature": 0
    }

    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {API_KEY}'
    }

    try:
        response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']
    except Exception as e:
        print(f"Error generating caption for {image_path}: {e}")
        return None

def process_images_in_folder(folder_path):
    """Process images in the specified folder."""
    for filename in os.listdir(folder_path):
        # Check if file is an image
        if any(filename.lower().endswith(ext) for ext in IMAGE_EXTENSIONS):
            image_path = os.path.join(folder_path, filename)
            txt_path = os.path.splitext(image_path)[0] + '.txt'

            # Check if txt file already exists
            if not os.path.exists(txt_path):
                print(f"Processing {filename}...")
                caption = generate_caption(image_path)

                if caption:
                    # Write caption to txt file
                    with open(txt_path, 'w', encoding='utf-8') as f:
                        f.write(caption)
                    print(f"Caption saved for {filename}")
                else:
                    print(f"Failed to generate caption for {filename}")

def main():
    # Check if API key is set
    if not API_KEY:
        print("Error: XAI_API_KEY environment variable not set.")
        return

    # Validate folder path
    if not os.path.exists(FOLDER_PATH):
        print(f"Error: Folder {FOLDER_PATH} does not exist.")
        return

    # Process images
    process_images_in_folder(FOLDER_PATH)

if __name__ == '__main__':
    main()
Edit Report
Pub: 08 Dec 2024 05:39 UTC
Edit: 08 Dec 2024 05:44 UTC
Views: 29