Add editing methods

.gitignore

@@ -0,0 +1 @@
+__pycache__

editor/operations/__init__.py

@@ -0,0 +1,3 @@
+from .add_noise import add_noise
+from .change_temperature import change_temperature
+from .add_random_colour_spill import add_random_colour_spill

editor/operations/add_noise.py

@@ -0,0 +1,11 @@
+import numpy as np
+from PIL import Image
+
+
+def add_noise(img: Image, alpha: float) -> Image:
+    img = img.convert("RGB")
+    width, height = img.size
+    random_colors = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
+    random_img = Image.fromarray(random_colors, mode="RGB")
+    result = Image.blend(img, random_img, alpha)
+    return result

editor/operations/add_random_colour_spill.py

@@ -0,0 +1,20 @@
+from PIL import Image
+from ..utils import random
+
+
+def add_random_colour_spill(image: Image, range: float) -> Image:
+    matrix = (
+        random(1 / range, range),
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        random(1 / range, range),
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        random(1 / range, range),
+        0.0,
+    )
+    return image.convert("RGB", matrix)

editor/operations/change_temperature.py

@@ -0,0 +1,42 @@
+from PIL import Image
+
+kelvin_table = {
+    1000: (255, 56, 0),
+    1500: (255, 109, 0),
+    2000: (255, 137, 18),
+    2500: (255, 161, 72),
+    3000: (255, 180, 107),
+    3500: (255, 196, 137),
+    4000: (255, 209, 163),
+    4500: (255, 219, 186),
+    5000: (255, 228, 206),
+    5500: (255, 236, 224),
+    6000: (255, 243, 239),
+    6500: (255, 249, 253),
+    7000: (245, 243, 255),
+    7500: (235, 238, 255),
+    8000: (227, 233, 255),
+    8500: (220, 229, 255),
+    9000: (214, 225, 255),
+    9500: (208, 222, 255),
+    10000: (204, 219, 255),
+}
+
+
+def change_temperature(image: Image, temperature: float) -> Image:
+    r, g, b = kelvin_table[temperature]
+    matrix = (
+        r / 255.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        g / 255.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        b / 255.0,
+        0.0,
+    )
+    return image.convert("RGB", matrix)

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)