diff --git a/sexp_effects/primitive_libs/color_ops.py b/sexp_effects/primitive_libs/color_ops.py index dd9076c..a0da497 100644 --- a/sexp_effects/primitive_libs/color_ops.py +++ b/sexp_effects/primitive_libs/color_ops.py @@ -8,14 +8,27 @@ import numpy as np import cv2 +def _to_numpy(img): + """Convert GPU frames or CuPy arrays to numpy for CPU processing.""" + # Handle GPUFrame objects + if hasattr(img, 'cpu'): + return img.cpu + # Handle CuPy arrays + if hasattr(img, 'get'): + return img.get() + return img + + def prim_adjust(img, brightness=0, contrast=1): """Adjust brightness and contrast. Brightness: -255 to 255, Contrast: 0 to 3+.""" + img = _to_numpy(img) result = (img.astype(np.float32) - 128) * contrast + 128 + brightness return np.clip(result, 0, 255).astype(np.uint8) -def prim_mix_gray(img, amount): +def prim_mix_gray(img_raw, amount): """Mix image with its grayscale version. 0=original, 1=grayscale.""" + img = _to_numpy(img_raw) gray = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2] gray_rgb = np.stack([gray, gray, gray], axis=-1) result = img.astype(np.float32) * (1 - amount) + gray_rgb * amount @@ -24,11 +37,13 @@ def prim_mix_gray(img, amount): def prim_invert_img(img): """Invert all pixel values.""" + img = _to_numpy(img) return (255 - img).astype(np.uint8) def prim_shift_hsv(img, h=0, s=1, v=1): """Shift HSV: h=degrees offset, s/v=multipliers.""" + img = _to_numpy(img) hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32) hsv[:, :, 0] = (hsv[:, :, 0] + h / 2) % 180 hsv[:, :, 1] = np.clip(hsv[:, :, 1] * s, 0, 255) @@ -38,6 +53,7 @@ def prim_shift_hsv(img, h=0, s=1, v=1): def prim_add_noise(img, amount): """Add gaussian noise to image.""" + img = _to_numpy(img) noise = np.random.normal(0, amount, img.shape) result = img.astype(np.float32) + noise return np.clip(result, 0, 255).astype(np.uint8) @@ -45,6 +61,7 @@ def prim_add_noise(img, amount): def prim_quantize(img, levels): """Reduce to N color levels per channel.""" + img = _to_numpy(img) levels = max(2, int(levels)) factor = 256 / levels result = (img // factor) * factor + factor // 2 @@ -53,6 +70,7 @@ def prim_quantize(img, levels): def prim_sepia(img, intensity=1.0): """Apply sepia tone effect.""" + img = _to_numpy(img) sepia_matrix = np.array([ [0.393, 0.769, 0.189], [0.349, 0.686, 0.168], @@ -65,6 +83,7 @@ def prim_sepia(img, intensity=1.0): def prim_grayscale(img): """Convert to grayscale (still RGB output).""" + img = _to_numpy(img) gray = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2] return np.stack([gray, gray, gray], axis=-1).astype(np.uint8)