test/effects/crt.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "opencv-python"]
# ///
"""
@effect crt
@version 1.0.0
@author artdag

@description
CRT / Scanlines effect. Simulates CRT monitor aesthetics with visible
scan lines, optional RGB subpixels, barrel distortion, and vignette.

@param line_spacing int
  @range 1 10
  @default 2
  Pixels between scanlines.

@param line_opacity float
  @range 0 1
  @default 0.3
  Darkness of scanlines.

@param rgb_subpixels bool
  @default false
  Show RGB subpixel pattern.

@param curvature float
  @range 0 0.5
  @default 0
  Barrel distortion amount for curved screen look.

@param vignette float
  @range 0 1
  @default 0
  Dark corners effect.

@param bloom float
  @range 0 1
  @default 0
  Glow/blur on bright areas.

@param flicker float
  @range 0 0.3
  @default 0
  Brightness variation.

@param seed int
  @default 42
  Random seed for flicker.

@state rng DeterministicRNG
  Random number generator for flicker.

@example
  (effect crt :line_spacing 2 :line_opacity 0.4)

@example
  ;; Full retro CRT look
  (effect crt :curvature 0.2 :vignette 0.3 :rgb_subpixels true :bloom 0.2)
"""

import numpy as np
import cv2
from pathlib import Path
import sys

# Import DeterministicRNG from same directory
_effects_dir = Path(__file__).parent
if str(_effects_dir) not in sys.path:
    sys.path.insert(0, str(_effects_dir))
from random import DeterministicRNG


def process_frame(frame: np.ndarray, params: dict, state: dict) -> tuple:
    """
    Apply CRT effect to a video frame.

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
        state: Persistent state dict

    Returns:
        Tuple of (processed_frame, new_state)
    """
    line_spacing = max(1, int(params.get("line_spacing", 2)))
    line_opacity = params.get("line_opacity", 0.3)
    rgb_subpixels = params.get("rgb_subpixels", False)
    curvature = params.get("curvature", 0)
    vignette = params.get("vignette", 0)
    bloom = params.get("bloom", 0)
    flicker = params.get("flicker", 0)
    seed = int(params.get("seed", 42))

    if state is None:
        state = {}

    # Initialize RNG
    if "rng" not in state:
        state["rng"] = DeterministicRNG(seed)
    rng = state["rng"]

    h, w = frame.shape[:2]
    result = frame.astype(np.float32).copy()

    # Apply barrel distortion (curvature)
    if curvature > 0:
        result = _apply_curvature(result, curvature)

    # Apply bloom (glow on bright areas)
    if bloom > 0:
        result = _apply_bloom(result, bloom)

    # Apply scanlines
    if line_opacity > 0:
        for y in range(0, h, line_spacing):
            result[y, :] = result[y, :] * (1 - line_opacity)

    # Apply RGB subpixel pattern
    if rgb_subpixels:
        for x in range(w):
            col_type = x % 3
            if col_type == 0:
                result[:, x, 0] *= 1.2
                result[:, x, 1] *= 0.8
                result[:, x, 2] *= 0.8
            elif col_type == 1:
                result[:, x, 0] *= 0.8
                result[:, x, 1] *= 1.2
                result[:, x, 2] *= 0.8
            else:
                result[:, x, 0] *= 0.8
                result[:, x, 1] *= 0.8
                result[:, x, 2] *= 1.2

    # Apply vignette
    if vignette > 0:
        y_coords, x_coords = np.ogrid[:h, :w]
        center_x, center_y = w / 2, h / 2
        dist = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
        max_dist = np.sqrt(center_x**2 + center_y**2)
        vignette_mask = 1 - (dist / max_dist) * vignette
        vignette_mask = np.clip(vignette_mask, 0, 1)
        result = result * vignette_mask[:, :, np.newaxis]

    # Apply flicker
    if flicker > 0:
        flicker_amount = 1.0 + rng.uniform(-flicker, flicker)
        result = result * flicker_amount

    return np.clip(result, 0, 255).astype(np.uint8), state


def _apply_curvature(frame: np.ndarray, strength: float) -> np.ndarray:
    """Apply barrel distortion."""
    h, w = frame.shape[:2]

    y_coords, x_coords = np.mgrid[0:h, 0:w].astype(np.float32)

    # Normalize to -1 to 1
    x_norm = (x_coords - w / 2) / (w / 2)
    y_norm = (y_coords - h / 2) / (h / 2)

    # Calculate radius
    r = np.sqrt(x_norm**2 + y_norm**2)

    # Apply barrel distortion
    r_distorted = r * (1 + strength * r**2)

    # Scale factor
    scale = np.where(r > 0, r_distorted / r, 1)

    # New coordinates
    new_x = (x_norm * scale * (w / 2) + w / 2).astype(np.float32)
    new_y = (y_norm * scale * (h / 2) + h / 2).astype(np.float32)

    result = cv2.remap(frame.astype(np.uint8), new_x, new_y,
                      cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT,
                      borderValue=(0, 0, 0))
    return result.astype(np.float32)


def _apply_bloom(frame: np.ndarray, strength: float) -> np.ndarray:
    """Apply bloom (glow on bright areas)."""
    gray = cv2.cvtColor(frame.astype(np.uint8), cv2.COLOR_RGB2GRAY)
    _, bright = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)

    bloom = cv2.GaussianBlur(bright, (21, 21), 0)
    bloom = cv2.cvtColor(bloom, cv2.COLOR_GRAY2RGB)

    result = frame + bloom.astype(np.float32) * strength * 0.5
    return result