config.py

"""Configuration for image dataset collection."""

import os
import sqlite3
from pathlib import Path

from dotenv import load_dotenv



# Load environment variables from .env file
load_dotenv(Path(__file__).parent / ".env")


def _get_env(key: str, default: str = "") -> str:
    """Get environment variable with optional default."""
    return os.environ.get(key, default)


# Image generation APIs
GROK_API_KEY = _get_env("GROK_API_KEY")

# AI Detection APIs
SIGHTENGINE_API_USER = _get_env("SIGHTENGINE_API_USER")
SIGHTENGINE_API_SECRET = _get_env("SIGHTENGINE_API_SECRET")
AIORNOT_API_KEY = _get_env("AIORNOT_API_KEY")
REALITYDEFENDER_API_KEY = _get_env("REALITYDEFENDER_API_KEY")
HIVE_API_KEY = _get_env("HIVE_API_KEY")
TRUTHSCAN_API_SECRET = _get_env("TRUTHSCAN_API_SECRET")
RESEMBLE_API_KEY = _get_env("RESEMBLE_API_KEY")
WINSTON_API_KEY = _get_env("WINSTON_API_KEY")
ILLUMINARTY_API_KEY = _get_env("ILLUMINARTY_API_KEY")
INAZA_API_KEY = _get_env("INAZA_API_KEY")
FAL_API_KEY = _get_env("FAL_API_KEY")
GEMINI_API_KEY = _get_env("GEMINI_API_KEY")
REPLICATE_API_TOKEN = _get_env("REPLICATE_API_TOKEN")

# LLM filtering
OPENAI_API_KEY = _get_env("OPENAI_API_KEY")
OPENAI_MODEL = "gpt-5.2"

# Storage - Local
DATA_DIR = "data"
DB_PATH = f"{DATA_DIR}/images.db"

# Storage - S3
S3_BUCKET = _get_env("S3_BUCKET")
S3_REGION = _get_env("S3_REGION", "us-west-1")

SEED = 972396464

# LLM validation prompts (tune these iteratively)
VALIDATION_PROMPTS = {
    "car_add_damage": """Analyze this image for use in a car damage dataset. The image must show a real, unedited photograph of a car exterior.

Reject the image if ANY of the following apply:
- Does not clearly show a car exterior (front, side, rear, or 3/4 view)
- Car fills less than 40% of the frame
- Less than 1/3 of the car is visible
- Image contains watermarks, text overlays, borders, or frames
- Image appears edited, composited, cut out, or digitally manipulated, including but not limited to:
  - unnaturally sharp or soft edges around the car
  - inconsistent lighting or shadows between the car and background
  - ground shadows that do not match the car's position
  - haloing, feathering, or edge artifacts
  - background blur or noise that does not match the car
  - car appears "pasted" onto the scene
- Is a render, CGI, illustration, stock composite, or dealership marketing image
- Poor quality (blurry, pixelated, over-processed)
- If there is reasonable suspicion of editing or compositing, reject the image.

Respond with JSON only:
{
  "valid": true/false,
  "reason": "brief reason if invalid, null if valid"
}""",
    "delivery_proof": """Analyze this image for use in a delivery proof dataset. The image should show a front door, porch, or doorstep area where a package could be placed.

Reject if ANY of these apply:
- Does not show a front door/porch/doorstep area
- **IMPORTANT**: Door must be within 14 feet of camera (reject if >14 feet away). Also reject if too close (<2 feet). This distance requirement is critical - the image should look like a typical delivery driver photo taken from arm's length to a few steps back.
- Has packages, people, or large obstructions in the door area
- Has watermarks, text overlays, borders, or frames
- Is clearly AI-generated (warped geometry, unnatural textures, impossible details)
- Is a render, illustration, or heavily edited image
- Poor quality or lighting that obscures details

Respond with JSON only:
{
    "valid": true/false,
    "reason": "brief reason if invalid, null if valid"
}""",
    "getty_editorial": """Analyze this image for use in a real photo dataset.

Reject if ANY of these apply:
- Has a stock photo watermark overlay (like "Getty Images", "Shutterstock", etc.) stamped on the image
- Image is blank, corrupted, or doesn't contain anything discernible
- Image is mostly solid color with no meaningful content

ACCEPT images that have:
- Natural scene text (street signs, building names, product labels, protest signs, name placards, etc.)
- Logos on clothing, vehicles, or buildings that are part of the scene
- Any other text that naturally appears in the photographed scene

The only text we reject is stock photo watermarks that are artificially overlaid on the image.

Respond with JSON only:
{
    "valid": true/false,
    "reason": "brief reason if invalid, null if valid"
}""",
    "product_add_defect": """Analyze this image for use in a product defect dataset. We need images of single products that could plausibly be modified to show a defect (scratch, crack, dent, stain, tear, missing part, etc.).

ACCEPT if ALL of these are true:
- Shows exactly ONE identifiable product as the main subject
- Product is a physical item that could have visible defects
- Real photograph (not illustration, render, or AI-generated)
- Product appears to be in good/new condition (no existing defects)

REJECT if ANY of these apply:
- Multiple products shown (we need a single focal product)
- Product is too small, distant, or out of focus
- Person is the main subject (product should be the focus, not a person wearing/holding it)
- Product already has visible damage, defects, or heavy wear
- Screenshot, collage, or edited/marked up image
- Blurry, dark, or very low quality image
- Product type where a visually obvious defect couldn't be added easily (liquids, powders, small consumables)

Respond with JSON only:
{
    "valid": true/false,
    "reason": "brief reason if invalid, null if valid"
}""",
    "food": """Analyze this image for use in a food quality dataset. We need clean images where the food could be modified to appear burnt, spoiled, or wrong.

ACCEPT if ALL are true:
- Shows food as the main subject (plated dish, ingredients, or prepared meal)
- Clear, decent quality photo (not blurry or dark)
- Food appears normal/appetizing (we will add defects later)
- Real photograph (not illustration or AI-generated)
- No overlays, text, stickers, borders, or editing artifacts

REJECT if ANY apply:
- No food visible or food is too small/distant
- Already shows burnt, spoiled, or defective food
- Person is the main subject (not the food)
- Has text overlays, watermarks, stickers, emojis, borders, or frames
- Screenshot, collage, or heavily edited/filtered image
- Very low quality, blurry, or dark

Respond with JSON only:
{
    "valid": true/false,
    "reason": "brief reason if invalid, null if valid"
}""",
}

# Alias for category name matching
VALIDATION_PROMPTS["product"] = VALIDATION_PROMPTS["product_add_defect"]


def _connect():
    """Connect to DB with truthscan variants normalized to 'truthscan'.

    Creates a temp view ``norm_det`` that merges truthscan / truthscan-noc2pa,
    preferring truthscan-noc2pa when both exist for the same image.
    """
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("""
        CREATE TEMP VIEW norm_det AS
        SELECT
            dr.image_ref_id,
            CASE WHEN dr.detector LIKE 'truthscan%' THEN 'truthscan'
                 ELSE dr.detector END as detector,
            dr.ai_score
        FROM detection_results dr
        WHERE NOT (dr.detector = 'truthscan'
            AND EXISTS (SELECT 1 FROM detection_results d2
                        WHERE d2.image_ref_id = dr.image_ref_id
                        AND d2.detector = 'truthscan-noc2pa'))
    """)
    return conn


# Model aliases: map variant names to a canonical name
MODEL_ALIASES = {
    "qwen-image-edit-2511": "qwen-image-2512",
    "seedream-v4.5-1k-t2i": "seedream-v4.5-1k",
}


def normalize_model(model: str) -> str:
    """Normalize model name using aliases.

    Also consolidates adversarial -t2i variants with their base model
    (e.g. adversarial-v4-grok-t2i -> adversarial-v4-grok).
    """
    model = MODEL_ALIASES.get(model, model)
    if model.startswith("adversarial-") and model.endswith("-t2i"):
        model = model[: -len("-t2i")]
    return model


# Display names for detectors (DB name -> LaTeX display name)
DETECTOR_DISPLAY_NAMES = {
    "aiornot": "AI or Not",
    "hive": "Hive AI",
    "illuminarty": "Illuminarty",
    "realitydefender": "RealityDefender",
    "resemble": "Resemble",
    "sightengine": "SightEngine",
    "truthscan": "TruthScan",
    "winston": "Winston AI",
}

# Display names for models (canonical DB name -> LaTeX display name)
MODEL_DISPLAY_NAMES = {
    "gemini-3-pro-image-preview": "Gemini",
    "gpt-image-1.5": "GPT Image 1.5",
    "grok-imagine-image-beta": "Grok",
    "qwen-image-2512": "Qwen",
    "seedream-v4.5-1k": "Seedream",
}

# Adversarial perturbation versions: internal DB key -> display name
# Ordered mild first, then strong
_ADV_VERSIONS = ["v7", "v4"]  # mild, strong
_VERSION_DISPLAY = {"v4": "Strong", "v7": "Mild"}

# Canonical order for the 5 main models
MAIN_MODELS = [
    "gemini-3-pro-image-preview",
    "gpt-image-1.5",
    "grok-imagine-image-beta",
    "qwen-image-2512",
    "seedream-v4.5-1k",
]