Add editing methods

editor/utils/__init__.py

@@ -0,0 +1,4 @@
+from .interpolate import interpolate
+from .random import random
+from .apply_pixel_shader import apply_pixel_shader
+from .get_colour_lut import get_colour_lut

editor/utils/apply_pixel_shader.py

@@ -0,0 +1,14 @@
+from typing import Callable, Tuple
+from PIL import Image
+
+
+def apply_pixel_shader(
+    img: Image, callback: Callable[[int, int, int], Tuple[int, int, int]]
+):
+    width, height = img.size
+    pixels = img.load()
+    for x in range(width):
+        for y in range(height):
+            r, g, b = pixels[x, y]
+            pixels[x, y] = callback(r, g, b)
+    return img

editor/utils/get_colour_lut.py

@@ -0,0 +1,21 @@
+import numpy as np
+from typing import List
+from .random import random
+from .interpolate import interpolate, INTERPOLATION_TYPE
+
+
+def get_edit_points(variance: float, count: int) -> List[float]:
+    return [
+        random(i / (count - 1) - variance, i / (count - 1) + variance)
+        for i in range(count)
+    ]
+
+
+def get_colour_lut(
+    variance=0.1, count=5, type: INTERPOLATION_TYPE = "cubic"
+) -> List[int]:
+    edit_points = get_edit_points(variance=variance, count=count)
+    return [
+        round(interpolate(edit_points, i / 255, type=type) * 255)
+        for i in np.linspace(0, 255, 256)
+    ]

editor/utils/interpolate.py

@@ -0,0 +1,35 @@
+import numpy as np
+from scipy.interpolate import CubicSpline
+from typing import List, Literal
+
+
+INTERPOLATION_TYPE = Literal["cubic", "linear"]
+
+
+def interpolate(
+    control_points: List[float], t: float, type: INTERPOLATION_TYPE
+) -> float:
+    control_points = sorted(control_points)
+
+    if type == "cubic":
+        x = np.linspace(0, 1, len(control_points))
+        cs = CubicSpline(x, control_points)
+        return cs(t)
+
+    if type == "linear":
+        n = len(control_points) - 1
+        segment_indices = np.linspace(0, 1, n + 1)
+
+        index = np.searchsorted(segment_indices, t, side="right") - 1
+
+        if t == 1:
+            return control_points[-1]
+        else:
+            t_normalized = (t - segment_indices[index]) / (
+                segment_indices[index + 1] - segment_indices[index]
+            )
+            return control_points[index] + t_normalized * (
+                control_points[index + 1] - control_points[index]
+            )
+
+    raise ValueError("Invalid type")

editor/utils/random.py

@@ -0,0 +1,10 @@
+import numpy as np
+
+
+def random(min: float = 0, max: float = 1):
+    mu = (max + min) / 2  # Mean of the distribution
+    sigma = (
+        max - min
+    ) / 6  # Standard deviation, chosen so that ~99.7% fall within [min_val, max_val]
+    sample = np.random.normal(mu, sigma)
+    return np.clip(sample, min, max)