From f5cf02865a9e59e0871c31aaf5202725b3e4294f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Slab=C3=BD?= <kron258@gmail.com>
Date: Sat, 6 Jan 2024 14:43:18 +0100
Subject: [PATCH] Add lightness and gamma control

---
 src/main.rs | 56 +++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 8896513..da762b6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,23 +40,24 @@ impl LinkyFrame {
         5 * 2
     }
 
-    fn set_pixel(&mut self, index: usize, rgb: (u8, u8, u8), brightness: f32) {
-        let r: f32 = f32::from(rgb.0) * brightness;
-        let g: f32 = f32::from(rgb.1) * brightness;
-        let b: f32 = f32::from(rgb.2) * brightness;
-        let rgb = unsafe {
+    fn set_pixel(&mut self, index: usize, rgb: (u8, u8, u8), brightness: f32, lightness: f32, gamma: f32) {
+        let r: f32 = (f32::from(rgb.0) * brightness / 255.0).powf(gamma) * 255.0;
+        let g: f32 = (f32::from(rgb.1) * brightness / 255.0).powf(gamma) * 255.0;
+        let b: f32 = (f32::from(rgb.2) * brightness / 255.0).powf(gamma) * 255.0;
+	let white = r.min(g).min(b) * lightness;
+        let rgbw = unsafe {
             (
-                r.to_int_unchecked(),
-                g.to_int_unchecked(),
-                b.to_int_unchecked(),
+                (r - white).to_int_unchecked(),
+                (g - white).to_int_unchecked(),
+                (b - white).to_int_unchecked(),
+		white.to_int_unchecked(),
             )
         };
 
-        let white = cmp::min(cmp::min(rgb.0, rgb.1), rgb.2);
-        self.data[index * 4] = rgb.0 - white;
-        self.data[index * 4 + 1] = rgb.1 - white;
-        self.data[index * 4 + 2] = rgb.2 - white;
-        self.data[index * 4 + 3] = white;
+        self.data[index * 4] = rgbw.0;
+        self.data[index * 4 + 1] = rgbw.1;
+        self.data[index * 4 + 2] = rgbw.2;
+        self.data[index * 4 + 3] = rgbw.3;
     }
 
     fn get_universe_data(&self, uni: u16) -> Vec<u8> {
@@ -273,7 +274,7 @@ fn blurred_pixel(data: &[(u8, u8, u8)], stride: usize, pix_index: usize, blur_ra
     )
 }
 
-fn gen_linky_frame(frame: &Video, index: u64, blur_size: u32, brightness: f32) -> LinkyFrame {
+fn gen_linky_frame(frame: &Video, index: u64, blur_size: u32, brightness: f32, lightness: f32, gamma: f32) -> LinkyFrame {
     let mut linky_frame = LinkyFrame::new();
     println!(
         "Processing frame {}: {} x {}",
@@ -308,7 +309,7 @@ fn gen_linky_frame(frame: &Video, index: u64, blur_size: u32, brightness: f32) -
                 video_plane[pix_index]
             };
 
-            linky_frame.set_pixel(i * LinkyFrame::rows() + j, rgb, brightness);
+            linky_frame.set_pixel(i * LinkyFrame::rows() + j, rgb, brightness, lightness, gamma);
         }
     }
     linky_frame
@@ -341,6 +342,18 @@ fn main() {
         "Normalized float scaling the brightness of each pixel (0.0-1.0 clamped)",
         "FLOAT",
     );
+    options.optopt(
+        "l",
+        "lightness",
+        "White level, used to scale the white channel value (0.0-1.0 clamped)",
+        "FLOAT",
+    );
+    options.optopt(
+        "g",
+        "gamma",
+        "Gamma exponent to use for color correction",
+        "FLOAT",
+    );
 
     let args: Vec<String> = env::args().collect();
     let program = args[0].clone();
@@ -367,6 +380,15 @@ fn main() {
         .unwrap_or_default()
         .unwrap_or(1.0);
     brightness = brightness.clamp(0.0, 1.0);
+    let mut lightness = matches
+        .opt_get::<f32>("l")
+        .unwrap_or_default()
+        .unwrap_or(1.0);
+    lightness = lightness.clamp(0.0, 1.0);
+    let gamma = matches
+        .opt_get::<f32>("g")
+        .unwrap_or_default()
+        .unwrap_or(1.0);
 
     if matches.opt_present("s") {
         let remote = matches.opt_str("s").unwrap();
@@ -378,7 +400,7 @@ fn main() {
 
         // let mut seq = 0;
         let frame_callback = |frame: &Video, index: u64| {
-            let linky_frame = gen_linky_frame(frame, index, blur_size, brightness);
+            let linky_frame = gen_linky_frame(frame, index, blur_size, brightness, lightness, gamma);
             for u in 0..LinkyFrame::universes() {
                 let mut packet = ArtNetDMX::new();
                 packet.set_universe(u);
@@ -407,7 +429,7 @@ fn main() {
     } else if matches.opt_present("w") {
         let mut file = File::create(matches.opt_str("w").unwrap()).unwrap();
         let frame_callback = |frame: &Video, index: u64| {
-            let linky_frame = gen_linky_frame(frame, index, blur_size, brightness);
+            let linky_frame = gen_linky_frame(frame, index, blur_size, brightness, lightness, gamma);
             for u in 0..LinkyFrame::universes() {
                 let mut packet = ArtNetDMX::new();
                 packet.set_universe(u);