Tune edits

src/editor/operations/get_colour_lut.py

@@ -16,11 +16,11 @@ def get_random_saturation_per_hue_lut(
 
 
 def get_random_brightness_lut(
-    variance: float = 0.2,
+    variance: float = 0.1,
     count: int = 6,
     type: INTERPOLATION_TYPE = "linear",
-    min_spectrum_size: float = 0.6,
-    max_spectrum_size: float = 1.15,
+    min_spectrum_size: float = 0.7,
+    max_spectrum_size: float = 1.1,
 ) -> List[int]:
     spectrum_size = np.random.uniform(min_spectrum_size, max_spectrum_size)
     spectrum_start = np.random.uniform(0, max(0, 1 - spectrum_size))
@@ -29,9 +29,9 @@ def get_random_brightness_lut(
             spectrum_start,
             *[
                 spectrum_start
-                + i * spectrum_size / (count - 2)
+                + i * spectrum_size / (count - 1)
                 + np.random.uniform(-variance, variance)
-                for i in range(1, count - 2)
+                for i in range(1, count - 1)
             ],
             spectrum_start + spectrum_size,
         ]

src/editor/training/random_edit.py

@@ -1,20 +1,32 @@
 from PIL import Image, ImageEnhance
-from ..utils import random, get_colour_lut, apply_pixel_shader
-from ..operations import add_noise, add_random_colour_spill
+from ..utils import random
+from ..operations import (
+    add_noise,
+    add_random_colour_spill,
+    get_random_gamma,
+    adjust_gamma,
+    apply_pixel_shader,
+    get_random_brightness_lut,
+    get_random_saturation_per_hue_lut,
+)
 import numpy as np
 
 
 def random_edit(img: Image, seed: int = 42) -> Image:
     np.random.seed(seed)
-    img = add_noise(img, random(0, 0.2))
-    img = ImageEnhance.Contrast(img).enhance(random(0.5, 2))
-    img = add_random_colour_spill(img, 1.3)
+    img = img.convert("RGB")
+
+    img = adjust_gamma(img, get_random_gamma())
+    img = add_noise(img, random(0, 0.1))
+    img = ImageEnhance.Contrast(img).enhance(random(0.5, 1.5))
+    img = add_random_colour_spill(img, 0.2)
 
     img = img.convert("HSV")
-    saturation_lut = get_colour_lut(variance=0.3, count=5, type="linear")
-    brightness_lut = get_colour_lut(variance=0.3, count=5, type="cubic")
+    saturation_lut = get_random_saturation_per_hue_lut()
+    brightness_lut = get_random_brightness_lut()
     img = apply_pixel_shader(
-        img, lambda h, s, v: (h, saturation_lut[s], brightness_lut[v])
+        img, lambda h, s, v: (h, round(s * saturation_lut[h]), brightness_lut[v])
     )
     img = img.convert("RGB")
+
     return img