EventSnap/e2e/loadtest/gen-images.py

#!/usr/bin/env python3
"""
Generate a pool of realistic phone-sized JPEGs for the EventSnap load test.

The stress test uploads ~1000 images, but they don't need to be 1000 unique
files — real load comes from realistic *byte size* and *decode cost*, which
drive bandwidth, the compression/preview pipeline (decode + 800x800 resize),
disk usage and the dynamic storage quota. So we generate a modest POOL of
distinct, high-entropy images (~2-4 MB, 12 MP, like a phone camera) and the
driver reuses them at random across the 1000 uploads.

High entropy matters: a flat gradient compresses to almost nothing and would
under-stress both the network and the JPEG decoder. We blend random noise over
a colorful gradient + big shapes so the files land in a realistic size band,
then tune JPEG quality per-image to hit the target size.

Output: $LT_PHOTOS_DIR (default /tmp/eventsnap-loadtest/photos)
Usage:  python3 gen-images.py [COUNT]     (default 40)
"""

import os
import sys
import math
import random
from PIL import Image, ImageDraw, ImageFont

OUT_DIR = os.environ.get("LT_PHOTOS_DIR", "/tmp/eventsnap-loadtest/photos")
COUNT = int(sys.argv[1]) if len(sys.argv) > 1 else 40

# Target JPEG size band (bytes). Typical modern phone photo.
TARGET_MIN = 2_000_000
TARGET_MAX = 4_500_000

# 12 MP-ish, both orientations (phones shoot portrait and landscape).
SIZES = [(4032, 3024), (3024, 4032)]

PALETTES = [
    [(250, 245, 235), (212, 175, 55), (120, 90, 30)],   # champagne / gold
    [(245, 244, 242), (190, 190, 198), (90, 92, 100)],   # silver / pearl
    [(255, 250, 240), (240, 200, 160), (170, 110, 70)],  # warm sunset
    [(235, 240, 248), (140, 170, 210), (40, 70, 120)],   # cool blue hour
    [(248, 240, 245), (210, 150, 180), (110, 50, 90)],   # rose dusk
    [(240, 248, 242), (150, 200, 170), (40, 110, 80)],   # garden green
]


def lerp(a, b, t):
    return tuple(int(a[i] + (b[i] - a[i]) * t) for i in range(3))


def gradient(w, h, palette, rng):
    """Diagonal 3-stop gradient base."""
    base = Image.new("RGB", (w, h))
    px = base.load()
    ang = rng.uniform(0, math.pi)
    dx, dy = math.cos(ang), math.sin(ang)
    # precompute per-column/row projection for speed
    maxproj = abs(dx) * w + abs(dy) * h
    for y in range(h):
        for x in range(0, w, 4):  # step 4 then fill — good enough, much faster
            t = (dx * x + dy * y) / maxproj
            t = min(1.0, max(0.0, t + rng.uniform(-0.02, 0.02)))
            if t < 0.5:
                c = lerp(palette[0], palette[1], t * 2)
            else:
                c = lerp(palette[1], palette[2], (t - 0.5) * 2)
            for k in range(4):
                if x + k < w:
                    px[x + k, y] = c
    return base


def add_shapes(img, rng):
    d = ImageDraw.Draw(img, "RGBA")
    w, h = img.size
    for _ in range(rng.randint(6, 14)):
        x0 = rng.randint(-w // 5, w)
        y0 = rng.randint(-h // 5, h)
        r = rng.randint(w // 12, w // 3)
        col = (rng.randint(0, 255), rng.randint(0, 255), rng.randint(0, 255), rng.randint(20, 90))
        if rng.random() < 0.5:
            d.ellipse([x0, y0, x0 + r, y0 + r], fill=col)
        else:
            d.rectangle([x0, y0, x0 + r, y0 + int(r * rng.uniform(0.4, 1.6))], fill=col)
    return img


def noisy(w, h, rng):
    """Full-resolution RGB noise from urandom — maximum entropy."""
    return Image.frombytes("RGB", (w, h), os.urandom(w * h * 3))


def label(img, idx, rng):
    d = ImageDraw.Draw(img)
    txt = f"EventSnap load #{idx:03d}"
    try:
        font = ImageFont.load_default(size=64)
    except TypeError:
        font = ImageFont.load_default()
    d.text((60, 60), txt, fill=(255, 255, 255), font=font)
    d.text((62, 62), txt, fill=(0, 0, 0), font=font)  # cheap shadow offset


def encode_to_band(img, path, rng):
    """Try qualities high→low until the file lands under TARGET_MAX; keep the
    first that also clears TARGET_MIN if possible."""
    best = None
    for q in (92, 88, 84, 80, 76, 72):
        img.save(path, "JPEG", quality=q, optimize=False)
        size = os.path.getsize(path)
        best = (q, size)
        if size <= TARGET_MAX:
            if size >= TARGET_MIN:
                return q, size
            # under the band — noise alpha likely too low; accept anyway at high q
            return q, size
    return best  # even q72 too big; accept the smallest we made


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    rng = random.Random(20260718)  # deterministic pool
    total_bytes = 0
    print(f"[gen] writing {COUNT} images to {OUT_DIR}")
    for i in range(COUNT):
        w, h = rng.choice(SIZES)
        palette = rng.choice(PALETTES)
        base = gradient(w, h, palette, rng)
        base = add_shapes(base, rng)
        # blend noise to inject entropy -> realistic JPEG size
        alpha = rng.uniform(0.28, 0.42)
        base = Image.blend(base, noisy(w, h, rng), alpha)
        label(base, i, rng)
        path = os.path.join(OUT_DIR, f"photo_{i:03d}.jpg")
        q, size = encode_to_band(base, path, rng)
        total_bytes += size
        print(f"  photo_{i:03d}.jpg  {w}x{h}  q{q}  {size/1_000_000:.2f} MB")
    avg = total_bytes / COUNT
    print(f"[gen] done. {COUNT} images, avg {avg/1_000_000:.2f} MB, "
          f"pool total {total_bytes/1_000_000:.1f} MB")
    print(f"[gen] projected for 1000 uploads (originals only): "
          f"~{avg*1000/1_000_000_000:.1f} GB")


if __name__ == "__main__":
    main()