image_processing.py

"""
Branchless Image Processing — Pixel manipulation without branches.

PATTERN: Arithmetic Masks for Pixel Operations
    In graphics programming (GPU shaders, SIMD intrinsics), branches are
    extremely expensive because they prevent the hardware from processing
    multiple pixels simultaneously (vectorization).

    Instead, we use arithmetic to compute results for ALL cases and then
    "mask" away the ones we don't want. The CPU does a tiny bit of extra
    math, but avoids the much more expensive branch misprediction penalty.

    All pixel values in these examples are integers in the range [0, 255].
"""


def branchless_clamp_byte(value):
    """
    Clamp an integer to [0, 255] without branching.

    This is the most common branchless operation in image processing —
    nearly every pixel calculation needs clamping to valid byte range.

    TRICK (two-step clamp):

    Step 1 — Clamp to minimum (0):
        below_zero = value >> 31      → -1 if negative, 0 if positive
        value = value & ~below_zero   → zeroed out if negative (AND with 0)

    Step 2 — Clamp to maximum (255):
        over_255 = (255 - value) >> 31  → -1 if value > 255, 0 otherwise
        value = value | (over_255 & 255)  → sets lower 8 bits to 1 if over
        value = value & 255               → keep only lowest 8 bits

    WORKED EXAMPLE (value = 300):
        Step 1: 300 >> 31 = 0, value stays 300
        Step 2: (255 - 300) >> 31 = -1 (negative → over)
                300 | (-1 & 255) = 300 | 255 = 0x12C | 0xFF = 0x1FF = 511
                511 & 255 = 255  ✓

    WORKED EXAMPLE (value = -50):
        Step 1: -50 >> 31 = -1
                ~(-1) = 0
                -50 & 0 = 0  ✓  (clamped to 0)
        Step 2: (255 - 0) >> 31 = 0 (positive → no change)
                0 & 255 = 0  ✓

    WORKED EXAMPLE (value = 128):
        Step 1: 128 >> 31 = 0, value stays 128
        Step 2: (255 - 128) = 127 >> 31 = 0 (positive → no change)
                128 & 255 = 128  ✓
    """
    # Clamp to min 0
    below_zero = value >> 31
    value = value & ~below_zero

    # Clamp to max 255
    over_255 = (255 - value) >> 31
    value = value | (over_255 & 255)
    value = value & 255

    return value


def branchless_threshold(pixel, threshold):
    """
    Binary threshold: returns 255 if pixel >= threshold, else 0.

    Used in: edge detection, document scanning, OCR preprocessing.

    TRICK:
        diff = pixel - threshold
        sign = diff >> 31          → 0 if pixel >= threshold, -1 if below
        result = 255 & ~sign       → 255 & ~0 = 255, or 255 & 0 = 0

    WHY ~sign WORKS:
        ~0  = -1 = 0xFFFFFFFF  (all 1s) → 255 & all-1s = 255
        ~(-1) = 0              (all 0s) → 255 & all-0s = 0

    WORKED EXAMPLE (pixel=200, threshold=128):
        diff = 200 - 128 = 72
        sign = 72 >> 31 = 0
        result = 255 & ~0 = 255 & 0xFFFFFFFF = 255  ✓

    WORKED EXAMPLE (pixel=50, threshold=128):
        diff = 50 - 128 = -78
        sign = -78 >> 31 = -1
        result = 255 & ~(-1) = 255 & 0 = 0  ✓
    """
    diff = pixel - threshold
    sign = diff >> 31
    return 255 & ~sign


def branchless_blend(pixel_a, pixel_b, alpha):
    """
    Alpha blend two pixels: result ≈ a × (alpha/255) + b × (1 − alpha/255)

    Used in: transparency, crossfades, compositing layers.

    We use 256 as the divisor instead of 255 because dividing by 256
    is just a right-shift by 8 (>> 8), which is much faster than actual
    division. The error is at most ±1, which is invisible in images.

    Formula: result = (a × alpha + b × (256 − alpha)) >> 8

    WORKED EXAMPLE (a=255, b=0, alpha=128 → ~50% blend):
        (255 × 128 + 0 × 128) >> 8
        = 32640 >> 8
        = 127  ✓  (halfway between 0 and 255)

    WORKED EXAMPLE (a=200, b=100, alpha=64 → ~25% blend):
        (200 × 64 + 100 × 192) >> 8
        = (12800 + 19200) >> 8
        = 32000 >> 8
        = 125  ✓  (closer to b=100, since alpha is low)
    """
    inv_alpha = 256 - alpha
    return (pixel_a * alpha + pixel_b * inv_alpha) >> 8


def branchless_contrast(pixel, factor):
    """
    Adjust contrast around midpoint (128).

    factor = 256 means "no change" (1.0×)
    factor > 256 increases contrast (e.g., 384 = 1.5×)
    factor < 256 decreases contrast (e.g., 128 = 0.5×)

    Formula:
        shifted = pixel - 128           (center around zero)
        adjusted = 128 + (shifted × factor) >> 8
        clamp to [0, 255]

    WORKED EXAMPLE (pixel=200, factor=384 = 1.5× contrast):
        shifted = 200 - 128 = 72
        adjusted = 128 + (72 × 384) >> 8 = 128 + 27648 >> 8 = 128 + 108 = 236
        clamp(236) = 236  ✓  (brighter pixels get pushed toward 255)

    WORKED EXAMPLE (pixel=50, factor=384 = 1.5× contrast):
        shifted = 50 - 128 = -78
        adjusted = 128 + (-78 × 384) >> 8 = 128 + (-29952) >> 8 = 128 - 117 = 11
        clamp(11) = 11  ✓  (darker pixels get pushed toward 0)
    """
    shifted = pixel - 128
    adjusted = 128 + ((shifted * factor) >> 8)
    return branchless_clamp_byte(adjusted)


def branchless_brightness(pixel, delta):
    """
    Adjust brightness by adding delta, clamped to [0, 255].

    This is the simplest image operation: just add/subtract and clamp.

    WORKED EXAMPLE (pixel=128, delta=200):
        128 + 200 = 328 → clamp → 255  ✓

    WORKED EXAMPLE (pixel=128, delta=-200):
        128 + (-200) = -72 → clamp → 0  ✓
    """
    return branchless_clamp_byte(pixel + delta)


def branchless_invert(pixel):
    """
    Invert pixel value: 255 - pixel.
    Trivially branchless — just arithmetic, no conditions needed.
    """
    return 255 - pixel


def branchless_grayscale(r, g, b):
    """
    Convert RGB to grayscale using the luminosity method.

    The human eye is most sensitive to green, then red, then blue.
    Standard weights: 0.299 R + 0.587 G + 0.114 B

    We approximate with integer math to avoid floating point:
        (r × 77 + g × 150 + b × 29) >> 8

    Verification: 77/256 ≈ 0.301, 150/256 ≈ 0.586, 29/256 ≈ 0.113  ✓

    WORKED EXAMPLE (pure red: r=255, g=0, b=0):
        (255 × 77 + 0 + 0) >> 8 = 19635 >> 8 = 76  ✓  (red appears dark)

    WORKED EXAMPLE (pure green: r=0, g=255, b=0):
        (0 + 255 × 150 + 0) >> 8 = 38250 >> 8 = 149  ✓  (green appears bright)
    """
    return (r * 77 + g * 150 + b * 29) >> 8


def branchless_select_channel(r, g, b, channel):
    """
    Select a color channel without branching.
    channel: 0=red, 1=green, 2=blue

    TRICK: Multiply each channel by a 0/1 flag and add them up.
    Only the selected channel's flag is 1; the others are 0.

    WORKED EXAMPLE (r=200, g=100, b=50, channel=1 → green):
        is_r = int(1 == 0) = 0
        is_g = int(1 == 1) = 1
        is_b = int(1 == 2) = 0
        result = 200 × 0 + 100 × 1 + 50 × 0 = 100  ✓
    """
    is_r = int(channel == 0)
    is_g = int(channel == 1)
    is_b = int(channel == 2)
    return r * is_r + g * is_g + b * is_b


def branchless_posterize(pixel, levels):
    """
    Reduce color depth to N levels (posterization) without branching.

    This creates the "poster art" effect by snapping pixel values to
    discrete steps.

    WORKED EXAMPLE (pixel=200, levels=4):
        step = 256 // 4 = 64
        (200 // 64) * 64 = 3 * 64 = 192  ✓  (snapped to nearest step)

    WORKED EXAMPLE (pixel=50, levels=4):
        (50 // 64) * 64 = 0 * 64 = 0  ✓
    """
    step = 256 // levels
    return (pixel // step) * step


# ──────────────────────────────────────────────
# Demo
# ──────────────────────────────────────────────

if __name__ == "__main__":
    print("Branchless Image Processing Demo")
    print("=" * 50)

    # Threshold
    print("\n--- Threshold (threshold=128) ---")
    print("  Pixels below 128 → 0 (black), above → 255 (white)")
    for p in [0, 64, 127, 128, 200, 255]:
        result = branchless_threshold(p, 128)
        print(f"  pixel {p:3d} → {result}")

    # Blend
    print("\n--- Alpha Blend (white=255, black=0) ---")
    print("  alpha=0 → all black, alpha=255 → all white")
    for alpha in [0, 64, 128, 192, 255]:
        result = branchless_blend(255, 0, alpha)
        print(f"  alpha={alpha:3d} → {result}")

    # Contrast
    print("\n--- Contrast Adjustment ---")
    for factor in [128, 256, 384]:
        label = f"{factor / 256:.1f}×"
        results = [branchless_contrast(p, factor) for p in [50, 128, 200]]
        print(f"  {label}: [50, 128, 200] → {results}")

    # Brightness
    print("\n--- Brightness Adjustment ---")
    for delta in [-50, 0, 50, 100]:
        result = branchless_brightness(128, delta)
        print(f"  128 + ({delta:+3d}) = {result}")

    # Grayscale
    print("\n--- Grayscale Conversion ---")
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (128, 128, 128)]
    color_names = ["Red", "Green", "Blue", "Gray"]
    for (r, g, b), name in zip(colors, color_names):
        gray = branchless_grayscale(r, g, b)
        print(f"  {name:5s} RGB({r:3d},{g:3d},{b:3d}) → Gray {gray}")

    # Posterize
    print("\n--- Posterize (4 levels) ---")
    for p in [0, 50, 100, 150, 200, 255]:
        result = branchless_posterize(p, 4)
        print(f"  {p:3d} → {result}")

    print()
    print("KEY TECHNIQUE: Arithmetic masks")
    print("  Replace if/else with multiply-by-0-or-1")
    print("  → enables SIMD vectorization across pixels")
    print("  → same math runs for every pixel, no branches")