token_estimator.py

#!/usr/bin/env python3
"""
Token Usage Estimator
Estimate Gemini API costs before processing
"""

import os
from pathlib import Path

# Image token calculation for Gemini
# Based on: https://ai.google.dev/gemini-api/docs/tokens

def estimate_image_tokens(width: int, height: int) -> int:
    """
    Estimate tokens for an image in Gemini API
    Images are resized and tiled internally
    """
    # Gemini 1.5: Images are charged at ~1280 tokens for standard images
    # Higher resolution images may cost more due to tiling

    # Base cost for image
    base_tokens = 1280

    # Additional cost for large images (auto-tiled)
    if width > 2048 or height > 2048:
        # Approximate tiling cost
        tiles = ((width + 1023) // 1024) * ((height + 1023) // 1024)
        base_tokens = base_tokens * max(1, tiles // 2)

    return base_tokens


def estimate_dataset_costs(
    rsicd_path: str,
    rsitmd_path: str,
    output_tokens_per_image: int = 150
):
    """Estimate costs for processing both datasets"""

    print("=" * 70)
    print("GEMINI API TOKEN USAGE ESTIMATION")
    print("=" * 70)

    # Count images
    rsicd_images = list_images(rsicd_path)
    rsitmd_images = list_images(rsitmd_path)

    total_images = len(rsicd_images) + len(rsitmd_images)

    print(f"\nDataset Analysis:")
    print(f"  RSICD:  {len(rsicd_images):,} images")
    print(f"  RSITMD: {len(rsitmd_images):,} images")
    print(f"  Total:  {total_images:,} images")

    # Sample image sizes for more accurate estimation
    print(f"\nImage Size Analysis:")

    avg_tokens = 1280  # Default for Gemini 1.5

    # Try to sample a few images
    sample_sizes = []
    for dataset_path, name in [(rsicd_path, "RSICD"), (rsitmd_path, "RSITMD")]:
        images = list_images(dataset_path)[:10]  # Sample first 10
        sizes = []
        for img_path in images:
            try:
                from PIL import Image
                with Image.open(img_path) as img:
                    sizes.append((img.width, img.height))
            except:
                pass

        if sizes:
            avg_w = sum(s[0] for s in sizes) // len(sizes)
            avg_h = sum(s[1] for s in sizes) // len(sizes)
            tokens = estimate_image_tokens(avg_w, avg_h)
            print(f"  {name}: avg size {avg_w}x{avg_h} → ~{tokens:,} input tokens/image")
            sample_sizes.append(tokens)

    if sample_sizes:
        avg_tokens = sum(sample_sizes) // len(sample_sizes)

    # Token calculations
    total_input_tokens = total_images * avg_tokens
    total_output_tokens = total_images * output_tokens_per_image
    total_tokens = total_input_tokens + total_output_tokens

    print(f"\nToken Usage Estimate:")
    print(f"  Input tokens:  {total_input_tokens:,}")
    print(f"  Output tokens: {total_output_tokens:,}")
    print(f"  Total tokens:  {total_tokens:,}")

    # Cost estimation for Gemini 1.5 Pro
    print(f"\nCost Estimation (Gemini 1.5 Pro):")
    input_cost_pro = total_input_tokens / 1_000_000 * 2.50
    output_cost_pro = total_output_tokens / 1_000_000 * 10.00
    total_cost_pro = input_cost_pro + output_cost_pro

    print(f"  Input:  ${input_cost_pro:.2f} (@ $2.50/1M tokens)")
    print(f"  Output: ${output_cost_pro:.2f} (@ $10.00/1M tokens)")
    print(f"  Total:  ${total_cost_pro:.2f}")

    # Cost estimation for Gemini 1.5 Flash (cheaper option)
    print(f"\nCost Estimation (Gemini 1.5 Flash - Recommended):")
    input_cost_flash = total_input_tokens / 1_000_000 * 0.075
    output_cost_flash = total_output_tokens / 1_000_000 * 0.30
    total_cost_flash = input_cost_flash + output_cost_flash

    print(f"  Input:  ${input_cost_flash:.2f} (@ $0.075/1M tokens)")
    print(f"  Output: ${output_cost_flash:.2f} (@ $0.30/1M tokens)")
    print(f"  Total:  ${total_cost_flash:.2f}")

    # Time estimation
    print(f"\nProcessing Time Estimate:")
    avg_time_per_image = 2  # seconds (conservative estimate)
    total_seconds = total_images * avg_time_per_image
    print(f"  At {avg_time_per_image}s per image: ~{total_seconds/3600:.1f} hours")

    print(f"\nRecommendations:")
    print(f"  • Use Gemini 1.5 Flash to save ${total_cost_pro - total_cost_flash:.2f}")
    print(f"  • Enable caching to avoid re-processing")
    print(f"  • Process in batches to allow checkpointing")
    print(f"  • Monitor token usage during processing")

    print("=" * 70)


def list_images(directory: str) -> list:
    """List all image files in directory"""
    extensions = ['.jpg', '.jpeg', '.png', '.tif', '.tiff', '.bmp']
    images = []
    path = Path(directory)

    if not path.exists():
        return images

    for ext in extensions:
        images.extend(path.glob(f"*{ext}"))
        images.extend(path.glob(f"*{ext.upper()}"))

    return images


def estimate_single_image(image_path: str):
    """Estimate tokens for a single image"""
    from PIL import Image

    with Image.open(image_path) as img:
        tokens = estimate_image_tokens(img.width, img.height)
        print(f"Image: {Path(image_path).name}")
        print(f"  Size: {img.width}x{img.height}")
        print(f"  Estimated input tokens: ~{tokens:,}")
        print(f"  Estimated output tokens: ~150 (typical)")
        print(f"  Total per image: ~{tokens + 150:,} tokens")


if __name__ == "__main__":
    from config import RSICD_IMAGE_PATH, RSITMD_IMAGE_PATH

    print("Estimating token costs for Beyond Pixels project\n")

    # Estimate both datasets
    estimate_dataset_costs(RSICD_IMAGE_PATH, RSITMD_IMAGE_PATH)

    print(f"\nTo estimate a single image:")
    print(f"  python token_estimator.py --image /path/to/image.jpg")