test/effects/pixelsort.py

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

@description
Pixel sorting glitch art effect. Sorts pixels within rows by brightness,
hue, or other properties. Creates distinctive streaked/melted aesthetics.

@param sort_by string
  @enum lightness hue saturation red green blue
  @default lightness
  Property to sort pixels by.

@param threshold_low float
  @range 0 255
  @default 50
  Pixels darker than this are not sorted.

@param threshold_high float
  @range 0 255
  @default 200
  Pixels brighter than this are not sorted.

@param angle float
  @range 0 180
  @default 0
  Sort direction: 0 = horizontal, 90 = vertical.

@param reverse bool
  @default false
  Reverse the sort order.

@example
  (effect pixelsort)

@example
  ;; Vertical pixel sort
  (effect pixelsort :angle 90)

@example
  ;; Sort by hue for rainbow streaks
  (effect pixelsort :sort_by "hue" :threshold_low 20 :threshold_high 240)
"""

import numpy as np
import cv2


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - sort_by: property to sort by (default "lightness")
            - threshold_low: min brightness to sort (default 50)
            - threshold_high: max brightness to sort (default 200)
            - angle: 0 = horizontal, 90 = vertical (default 0)
            - reverse: reverse sort order (default False)
        state: Persistent state dict (unused)

    Returns:
        Tuple of (processed_frame, new_state)
    """
    sort_by = params.get("sort_by", "lightness")
    threshold_low = params.get("threshold_low", 50)
    threshold_high = params.get("threshold_high", 200)
    angle = params.get("angle", 0)
    reverse = params.get("reverse", False)

    h, w = frame.shape[:2]

    # Rotate for non-horizontal sorting
    if 45 <= (angle % 180) <= 135:
        frame = np.transpose(frame, (1, 0, 2))
        h, w = frame.shape[:2]
        rotated = True
    else:
        rotated = False

    result = frame.copy()

    # Get sort values
    sort_values = _get_sort_values(frame, sort_by)

    # Create mask of pixels to sort
    mask = (sort_values >= threshold_low) & (sort_values <= threshold_high)

    # Sort each row
    for y in range(h):
        row = result[y].copy()
        row_mask = mask[y]
        row_values = sort_values[y]

        # Find contiguous segments to sort
        segments = _find_segments(row_mask)

        for start, end in segments:
            if end - start > 1:
                segment_values = row_values[start:end]
                sort_indices = np.argsort(segment_values)
                if reverse:
                    sort_indices = sort_indices[::-1]
                row[start:end] = row[start:end][sort_indices]

        result[y] = row

    # Rotate back if needed
    if rotated:
        result = np.transpose(result, (1, 0, 2))

    return np.ascontiguousarray(result), state


def _get_sort_values(frame, sort_by):
    """Get values to sort pixels by."""
    if sort_by == "lightness":
        return cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).astype(np.float32)
    elif sort_by == "hue":
        hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
        return hsv[:, :, 0].astype(np.float32)
    elif sort_by == "saturation":
        hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
        return hsv[:, :, 1].astype(np.float32)
    elif sort_by == "red":
        return frame[:, :, 0].astype(np.float32)
    elif sort_by == "green":
        return frame[:, :, 1].astype(np.float32)
    elif sort_by == "blue":
        return frame[:, :, 2].astype(np.float32)
    return cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).astype(np.float32)


def _find_segments(mask):
    """Find contiguous True segments in mask."""
    segments = []
    start = None
    for i, val in enumerate(mask):
        if val and start is None:
            start = i
        elif not val and start is not None:
            segments.append((start, i))
            start = None
    if start is not None:
        segments.append((start, len(mask)))
    return segments