iris/color/
mod.rs

1use crate::error::{IrisError, Result};
2use crate::image::Image;
3use burn::tensor::{Tensor, TensorData, backend::Backend};
4
5/// Color space conversion utilities for images.
6impl<B: Backend> Image<B> {
7    /// Converts an RGB image to HSV (Hue, Saturation, Value) color space.
8    /// Input must be a 3-channel image with values in [0, 1].
9    /// Returns a 3-channel image where H is in [0, 360]/360 (normalized to [0, 1]),
10    /// S is in [0, 1], and V is in [0, 1].
11    pub fn rgb_to_hsv(&self) -> Result<Self> {
12        let dims = self.tensor.dims();
13        let c = dims[0];
14        let h = dims[1];
15        let w = dims[2];
16
17        if c != 3 {
18            return Err(IrisError::InvalidParameter(
19                "Input must be a 3-channel RGB image".into(),
20            ));
21        }
22
23        let tensor_data = self.tensor.clone().into_data();
24        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
25        let mut out_vals = vec![0.0f32; 3 * h * w];
26
27        let pixels = h * w;
28
29        for i in 0..pixels {
30            let r = flat_vals[i];
31            let g = flat_vals[pixels + i];
32            let b = flat_vals[2 * pixels + i];
33
34            let max = r.max(g).max(b);
35            let min = r.min(g).min(b);
36            let delta = max - min;
37
38            // Value
39            out_vals[2 * pixels + i] = max;
40
41            // Saturation
42            out_vals[pixels + i] = if max.abs() < 1e-6 { 0.0 } else { delta / max };
43
44            // Hue
45            let hue = if delta.abs() < 1e-6 {
46                0.0
47            } else if (max - r).abs() < 1e-6 {
48                60.0 * (((g - b) / delta) % 6.0)
49            } else if (max - g).abs() < 1e-6 {
50                60.0 * (((b - r) / delta) + 2.0)
51            } else {
52                60.0 * (((r - g) / delta) + 4.0)
53            };
54
55            // Normalize hue to [0, 1]
56            let hue_norm = if hue < 0.0 {
57                (hue + 360.0) / 360.0
58            } else {
59                hue / 360.0
60            };
61            out_vals[i] = hue_norm;
62        }
63
64        let device = self.tensor.device();
65        let data = TensorData::new(out_vals, [3, h, w]);
66        let tensor = Tensor::<B, 3>::from_data(data, &device);
67        Ok(Image::new(tensor))
68    }
69
70    /// Converts an HSV image to RGB color space.
71    /// Input must be a 3-channel image where H is in [0, 1] (normalized from 360),
72    /// S is in [0, 1], and V is in [0, 1].
73    pub fn hsv_to_rgb(&self) -> Result<Self> {
74        let dims = self.tensor.dims();
75        let c = dims[0];
76        let h_dim = dims[1];
77        let w = dims[2];
78
79        if c != 3 {
80            return Err(IrisError::InvalidParameter(
81                "Input must be a 3-channel HSV image".into(),
82            ));
83        }
84
85        let tensor_data = self.tensor.clone().into_data();
86        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
87        let mut out_vals = vec![0.0f32; 3 * h_dim * w];
88
89        let pixels = h_dim * w;
90
91        for i in 0..pixels {
92            let hue = flat_vals[i] * 360.0; // Denormalize hue
93            let sat = flat_vals[pixels + i];
94            let val = flat_vals[2 * pixels + i];
95
96            let c_val = val * sat;
97            let x = c_val * (1.0 - ((hue / 60.0) % 2.0 - 1.0).abs());
98            let m = val - c_val;
99
100            let (r, g, b) = if hue < 60.0 {
101                (c_val, x, 0.0)
102            } else if hue < 120.0 {
103                (x, c_val, 0.0)
104            } else if hue < 180.0 {
105                (0.0, c_val, x)
106            } else if hue < 240.0 {
107                (0.0, x, c_val)
108            } else if hue < 300.0 {
109                (x, 0.0, c_val)
110            } else {
111                (c_val, 0.0, x)
112            };
113
114            out_vals[i] = r + m;
115            out_vals[pixels + i] = g + m;
116            out_vals[2 * pixels + i] = b + m;
117        }
118
119        let device = self.tensor.device();
120        let data = TensorData::new(out_vals, [3, h_dim, w]);
121        let tensor = Tensor::<B, 3>::from_data(data, &device);
122        Ok(Image::new(tensor))
123    }
124
125    /// Converts an RGB image to HLS (Hue, Lightness, Saturation) color space.
126    /// H is normalized to [0, 1] (from 360 degrees), L and S are in [0, 1].
127    pub fn rgb_to_hls(&self) -> Result<Self> {
128        let dims = self.tensor.dims();
129        let c = dims[0];
130        let h = dims[1];
131        let w = dims[2];
132
133        if c != 3 {
134            return Err(IrisError::InvalidParameter(
135                "Input must be a 3-channel RGB image".into(),
136            ));
137        }
138
139        let tensor_data = self.tensor.clone().into_data();
140        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
141        let mut out_vals = vec![0.0f32; 3 * h * w];
142
143        let pixels = h * w;
144
145        for i in 0..pixels {
146            let r = flat_vals[i];
147            let g = flat_vals[pixels + i];
148            let b = flat_vals[2 * pixels + i];
149
150            let max = r.max(g).max(b);
151            let min = r.min(g).min(b);
152            let delta = max - min;
153
154            // Lightness
155            let l = (max + min) / 2.0;
156            out_vals[pixels + i] = l;
157
158            // Saturation
159            out_vals[2 * pixels + i] = if delta.abs() < 1e-6 {
160                0.0
161            } else if l < 0.5 {
162                delta / (max + min)
163            } else {
164                delta / (2.0 - max - min)
165            };
166
167            // Hue
168            let hue = if delta.abs() < 1e-6 {
169                0.0
170            } else if (max - r).abs() < 1e-6 {
171                60.0 * (((g - b) / delta) % 6.0)
172            } else if (max - g).abs() < 1e-6 {
173                60.0 * (((b - r) / delta) + 2.0)
174            } else {
175                60.0 * (((r - g) / delta) + 4.0)
176            };
177
178            let hue_norm = if hue < 0.0 {
179                (hue + 360.0) / 360.0
180            } else {
181                hue / 360.0
182            };
183            out_vals[i] = hue_norm;
184        }
185
186        let device = self.tensor.device();
187        let data = TensorData::new(out_vals, [3, h, w]);
188        let tensor = Tensor::<B, 3>::from_data(data, &device);
189        Ok(Image::new(tensor))
190    }
191
192    /// Converts an HLS image to RGB color space.
193    pub fn hls_to_rgb(&self) -> Result<Self> {
194        let dims = self.tensor.dims();
195        let c = dims[0];
196        let h_dim = dims[1];
197        let w = dims[2];
198
199        if c != 3 {
200            return Err(IrisError::InvalidParameter(
201                "Input must be a 3-channel HLS image".into(),
202            ));
203        }
204
205        let tensor_data = self.tensor.clone().into_data();
206        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
207        let mut out_vals = vec![0.0f32; 3 * h_dim * w];
208
209        let pixels = h_dim * w;
210
211        for i in 0..pixels {
212            let hue = flat_vals[i] * 360.0;
213            let l = flat_vals[pixels + i];
214            let s = flat_vals[2 * pixels + i];
215
216            let c_val = (1.0 - (2.0 * l - 1.0).abs()) * s;
217            let x = c_val * (1.0 - ((hue / 60.0) % 2.0 - 1.0).abs());
218            let m = l - c_val / 2.0;
219
220            let (r, g, b) = if hue < 60.0 {
221                (c_val, x, 0.0)
222            } else if hue < 120.0 {
223                (x, c_val, 0.0)
224            } else if hue < 180.0 {
225                (0.0, c_val, x)
226            } else if hue < 240.0 {
227                (0.0, x, c_val)
228            } else if hue < 300.0 {
229                (x, 0.0, c_val)
230            } else {
231                (c_val, 0.0, x)
232            };
233
234            out_vals[i] = r + m;
235            out_vals[pixels + i] = g + m;
236            out_vals[2 * pixels + i] = b + m;
237        }
238
239        let device = self.tensor.device();
240        let data = TensorData::new(out_vals, [3, h_dim, w]);
241        let tensor = Tensor::<B, 3>::from_data(data, &device);
242        Ok(Image::new(tensor))
243    }
244
245    /// Splits a multi-channel image into individual single-channel images.
246    pub fn split_channels(&self) -> Result<Vec<Self>> {
247        let dims = self.tensor.dims();
248        let c = dims[0];
249        let h = dims[1];
250        let w = dims[2];
251
252        let tensor_data = self.tensor.clone().into_data();
253        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
254        let pixels = h * w;
255
256        let mut channels = Vec::with_capacity(c);
257        for ch in 0..c {
258            let start = ch * pixels;
259            let channel_data = flat_vals[start..start + pixels].to_vec();
260            let data = TensorData::new(channel_data, [1, h, w]);
261            let tensor = Tensor::<B, 3>::from_data(data, &self.tensor.device());
262            channels.push(Image::new(tensor));
263        }
264
265        Ok(channels)
266    }
267
268    /// Merges single-channel images into a multi-channel image.
269    pub fn merge_channels(channels: &[Image<B>]) -> Result<Self> {
270        if channels.is_empty() {
271            return Err(IrisError::InvalidParameter(
272                "At least one channel is required".into(),
273            ));
274        }
275
276        let dims = channels[0].tensor.dims();
277        let h = dims[1];
278        let w = dims[2];
279        let c = channels.len();
280
281        let mut all_vals = Vec::with_capacity(c * h * w);
282        for ch in channels {
283            let ch_dims = ch.tensor.dims();
284            if ch_dims[1] != h || ch_dims[2] != w {
285                return Err(IrisError::DimensionMismatch {
286                    expected: vec![1, h, w],
287                    actual: vec![ch_dims[0], ch_dims[1], ch_dims[2]],
288                });
289            }
290            let data = ch.tensor.clone().into_data();
291            let vals: Vec<f32> = data.iter::<f32>().collect();
292            all_vals.extend_from_slice(&vals);
293        }
294
295        let device = channels[0].tensor.device();
296        let data = TensorData::new(all_vals, [c, h, w]);
297        let tensor = Tensor::<B, 3>::from_data(data, &device);
298        Ok(Image::new(tensor))
299    }
300
301    /// Converts RGB to CIE XYZ color space.
302    /// Uses sRGB D65 illuminant matrix.
303    pub fn rgb_to_xyz(&self) -> Result<Self> {
304        let dims = self.tensor.dims();
305        if dims[0] != 3 {
306            return Err(IrisError::InvalidParameter(
307                "Input must be 3-channel RGB".into(),
308            ));
309        }
310        let h = dims[1];
311        let w = dims[2];
312        let data = self.tensor.clone().into_data();
313        let flat: Vec<f32> = data.iter::<f32>().collect();
314        let pixels = h * w;
315        let mut out = vec![0.0f32; 3 * pixels];
316
317        for i in 0..pixels {
318            // sRGB gamma linearization
319            let r_lin = linearize(flat[i]);
320            let g_lin = linearize(flat[pixels + i]);
321            let b_lin = linearize(flat[2 * pixels + i]);
322
323            // sRGB to XYZ (D65)
324            out[i] = 0.412_456_4 * r_lin + 0.357_576_1 * g_lin + 0.180_437_5 * b_lin;
325            out[pixels + i] = 0.212_672_9 * r_lin + 0.715_152_2 * g_lin + 0.072_175_0 * b_lin;
326            out[2 * pixels + i] = 0.019_333_9 * r_lin + 0.119_192 * g_lin + 0.950_304_1 * b_lin;
327        }
328
329        Ok(Image::new(Tensor::<B, 3>::from_data(
330            TensorData::new(out, [3, h, w]),
331            &self.tensor.device(),
332        )))
333    }
334
335    /// Converts CIE XYZ to RGB color space.
336    pub fn xyz_to_rgb(&self) -> Result<Self> {
337        let dims = self.tensor.dims();
338        if dims[0] != 3 {
339            return Err(IrisError::InvalidParameter(
340                "Input must be 3-channel XYZ".into(),
341            ));
342        }
343        let h = dims[1];
344        let w = dims[2];
345        let data = self.tensor.clone().into_data();
346        let flat: Vec<f32> = data.iter::<f32>().collect();
347        let pixels = h * w;
348        let mut out = vec![0.0f32; 3 * pixels];
349
350        for i in 0..pixels {
351            let x = flat[i];
352            let y = flat[pixels + i];
353            let z = flat[2 * pixels + i];
354
355            // XYZ to linear sRGB
356            let r_lin = 3.240_454_2 * x - 1.537_138_5 * y - 0.498_531_4 * z;
357            let g_lin = -0.969_266 * x + 1.876_010_8 * y + 0.041_556_0 * z;
358            let b_lin = 0.055_643_4 * x - 0.204_025_9 * y + 1.057_225_2 * z;
359
360            // Gamma encoding
361            out[i] = delinearize(r_lin);
362            out[pixels + i] = delinearize(g_lin);
363            out[2 * pixels + i] = delinearize(b_lin);
364        }
365
366        Ok(Image::new(Tensor::<B, 3>::from_data(
367            TensorData::new(out, [3, h, w]),
368            &self.tensor.device(),
369        )))
370    }
371
372    /// Converts RGB to CIE L*a*b* color space.
373    /// Pipeline: RGB -> XYZ -> L*a*b* (D65 white point).
374    pub fn rgb_to_lab(&self) -> Result<Self> {
375        let xyz = self.rgb_to_xyz()?;
376        let dims = xyz.tensor.dims();
377        let h = dims[1];
378        let w = dims[2];
379        let data = xyz.tensor.clone().into_data();
380        let flat: Vec<f32> = data.iter::<f32>().collect();
381        let pixels = h * w;
382        let mut out = vec![0.0f32; 3 * pixels];
383
384        // D65 white point
385        let xn = 0.950_47_f64;
386        let yn = 1.0_f64;
387        let zn = 1.088_83_f64;
388
389        for i in 0..pixels {
390            let x = flat[i] as f64 / xn;
391            let y = flat[pixels + i] as f64 / yn;
392            let z = flat[2 * pixels + i] as f64 / zn;
393
394            let fx = lab_f(x);
395            let fy = lab_f(y);
396            let fz = lab_f(z);
397
398            let l = 116.0 * fy - 16.0;
399            let a = 500.0 * (fx - fy);
400            let b = 200.0 * (fy - fz);
401
402            out[i] = (l / 100.0) as f32; // L in [0, 1]
403            out[pixels + i] = ((a + 128.0) / 255.0) as f32; // a in [0, 1]
404            out[2 * pixels + i] = ((b + 128.0) / 255.0) as f32; // b in [0, 1]
405        }
406
407        Ok(Image::new(Tensor::<B, 3>::from_data(
408            TensorData::new(out, [3, h, w]),
409            &self.tensor.device(),
410        )))
411    }
412
413    /// Converts CIE L*a*b* to RGB color space.
414    pub fn lab_to_rgb(&self) -> Result<Self> {
415        let dims = self.tensor.dims();
416        if dims[0] != 3 {
417            return Err(IrisError::InvalidParameter(
418                "Input must be 3-channel LAB".into(),
419            ));
420        }
421        let h = dims[1];
422        let w = dims[2];
423        let data = self.tensor.clone().into_data();
424        let flat: Vec<f32> = data.iter::<f32>().collect();
425        let pixels = h * w;
426        let mut xyz_vals = vec![0.0f32; 3 * pixels];
427
428        let xn = 0.950_47_f64;
429        let yn = 1.0_f64;
430        let zn = 1.088_83_f64;
431
432        for i in 0..pixels {
433            let l = flat[i] as f64 * 100.0;
434            let a = flat[pixels + i] as f64 * 255.0 - 128.0;
435            let b = flat[2 * pixels + i] as f64 * 255.0 - 128.0;
436
437            let fy = (l + 16.0) / 116.0;
438            let fx = a / 500.0 + fy;
439            let fz = fy - b / 200.0;
440
441            let x = lab_f_inv(fx) * xn;
442            let y = lab_f_inv(fy) * yn;
443            let z = lab_f_inv(fz) * zn;
444
445            xyz_vals[i] = x as f32;
446            xyz_vals[pixels + i] = y as f32;
447            xyz_vals[2 * pixels + i] = z as f32;
448        }
449
450        let xyz_img = Image::new(Tensor::<B, 3>::from_data(
451            TensorData::new(xyz_vals, [3, h, w]),
452            &self.tensor.device(),
453        ));
454        xyz_img.xyz_to_rgb()
455    }
456
457    /// Converts RGB to YUV (BT.601) color space.
458    pub fn rgb_to_yuv(&self) -> Result<Self> {
459        let dims = self.tensor.dims();
460        if dims[0] != 3 {
461            return Err(IrisError::InvalidParameter(
462                "Input must be 3-channel RGB".into(),
463            ));
464        }
465        let h = dims[1];
466        let w = dims[2];
467        let data = self.tensor.clone().into_data();
468        let flat: Vec<f32> = data.iter::<f32>().collect();
469        let pixels = h * w;
470        let mut out = vec![0.0f32; 3 * pixels];
471
472        for i in 0..pixels {
473            let r = flat[i] as f64;
474            let g = flat[pixels + i] as f64;
475            let b = flat[2 * pixels + i] as f64;
476
477            let y = 0.299 * r + 0.587 * g + 0.114 * b;
478            let u = -0.147_13 * r - 0.288_86 * g + 0.436 * b + 0.5;
479            let v = 0.615 * r - 0.514_99 * g - 0.100_01 * b + 0.5;
480
481            out[i] = y.clamp(0.0, 1.0) as f32;
482            out[pixels + i] = u.clamp(0.0, 1.0) as f32;
483            out[2 * pixels + i] = v.clamp(0.0, 1.0) as f32;
484        }
485
486        Ok(Image::new(Tensor::<B, 3>::from_data(
487            TensorData::new(out, [3, h, w]),
488            &self.tensor.device(),
489        )))
490    }
491
492    /// Converts YUV (BT.601) to RGB color space.
493    pub fn yuv_to_rgb(&self) -> Result<Self> {
494        let dims = self.tensor.dims();
495        if dims[0] != 3 {
496            return Err(IrisError::InvalidParameter(
497                "Input must be 3-channel YUV".into(),
498            ));
499        }
500        let h = dims[1];
501        let w = dims[2];
502        let data = self.tensor.clone().into_data();
503        let flat: Vec<f32> = data.iter::<f32>().collect();
504        let pixels = h * w;
505        let mut out = vec![0.0f32; 3 * pixels];
506
507        for i in 0..pixels {
508            let y = flat[i] as f64;
509            let u = flat[pixels + i] as f64 - 0.5;
510            let v = flat[2 * pixels + i] as f64 - 0.5;
511
512            let r = y + 1.139_83 * v;
513            let g = y - 0.394_65 * u - 0.580_60 * v;
514            let b = y + 2.032_11 * u;
515
516            out[i] = r.clamp(0.0, 1.0) as f32;
517            out[pixels + i] = g.clamp(0.0, 1.0) as f32;
518            out[2 * pixels + i] = b.clamp(0.0, 1.0) as f32;
519        }
520
521        Ok(Image::new(Tensor::<B, 3>::from_data(
522            TensorData::new(out, [3, h, w]),
523            &self.tensor.device(),
524        )))
525    }
526
527    /// Converts RGB to YCrCb (BT.601) color space.
528    pub fn rgb_to_ycrcb(&self) -> Result<Self> {
529        let dims = self.tensor.dims();
530        if dims[0] != 3 {
531            return Err(IrisError::InvalidParameter(
532                "Input must be 3-channel RGB".into(),
533            ));
534        }
535        let h = dims[1];
536        let w = dims[2];
537        let data = self.tensor.clone().into_data();
538        let flat: Vec<f32> = data.iter::<f32>().collect();
539        let pixels = h * w;
540        let mut out = vec![0.0f32; 3 * pixels];
541
542        for i in 0..pixels {
543            let r = flat[i] as f64;
544            let g = flat[pixels + i] as f64;
545            let b = flat[2 * pixels + i] as f64;
546
547            let y = 0.299 * r + 0.587 * g + 0.114 * b;
548            let cr = 0.713 * (r - y) + 0.5;
549            let cb = 0.564 * (b - y) + 0.5;
550
551            out[i] = y.clamp(0.0, 1.0) as f32;
552            out[pixels + i] = cr.clamp(0.0, 1.0) as f32;
553            out[2 * pixels + i] = cb.clamp(0.0, 1.0) as f32;
554        }
555
556        Ok(Image::new(Tensor::<B, 3>::from_data(
557            TensorData::new(out, [3, h, w]),
558            &self.tensor.device(),
559        )))
560    }
561
562    /// Converts an RGB image to CMYK (Cyan, Magenta, Yellow, Key/Black) color space.
563    /// Input must be a 3-channel image with values in [0, 1].
564    /// Returns a 4-channel image with values in [0, 1].
565    pub fn rgb_to_cmyk(&self) -> Result<Self> {
566        let dims = self.tensor.dims();
567        if dims[0] != 3 {
568            return Err(IrisError::InvalidParameter(
569                "Input must be a 3-channel RGB image".into(),
570            ));
571        }
572        let h = dims[1];
573        let w = dims[2];
574        let pixels = h * w;
575
576        let data = self.tensor.clone().into_data();
577        let flat: Vec<f32> = data.iter::<f32>().collect();
578        let mut out = vec![0.0f32; 4 * pixels];
579
580        for i in 0..pixels {
581            let r = flat[i];
582            let g = flat[pixels + i];
583            let b = flat[2 * pixels + i];
584
585            let k = 1.0f32 - r.max(g).max(b);
586            if k < 1.0 - 1e-6 {
587                let inv = 1.0 / (1.0 - k);
588                out[i] = (1.0 - r - k) * inv; // Cyan
589                out[pixels + i] = (1.0 - g - k) * inv; // Magenta
590                out[2 * pixels + i] = (1.0 - b - k) * inv; // Yellow
591            } else {
592                out[i] = 0.0;
593                out[pixels + i] = 0.0;
594                out[2 * pixels + i] = 0.0;
595            }
596            out[3 * pixels + i] = k; // Black
597        }
598
599        let device = self.tensor.device();
600        let tensor = Tensor::<B, 3>::from_data(TensorData::new(out, [4, h, w]), &device);
601        Ok(Image::new(tensor))
602    }
603
604    /// Converts a CMYK image to RGB color space.
605    /// Input must be a 4-channel image with values in [0, 1].
606    /// Returns a 3-channel RGB image with values in [0, 1].
607    pub fn cmyk_to_rgb(&self) -> Result<Self> {
608        let dims = self.tensor.dims();
609        if dims[0] != 4 {
610            return Err(IrisError::InvalidParameter(
611                "Input must be a 4-channel CMYK image".into(),
612            ));
613        }
614        let h = dims[1];
615        let w = dims[2];
616        let pixels = h * w;
617
618        let data = self.tensor.clone().into_data();
619        let flat: Vec<f32> = data.iter::<f32>().collect();
620        let mut out = vec![0.0f32; 3 * pixels];
621
622        for i in 0..pixels {
623            let c = flat[i];
624            let m = flat[pixels + i];
625            let y = flat[2 * pixels + i];
626            let k = flat[3 * pixels + i];
627
628            out[i] = (1.0 - c) * (1.0 - k); // Red
629            out[pixels + i] = (1.0 - m) * (1.0 - k); // Green
630            out[2 * pixels + i] = (1.0 - y) * (1.0 - k); // Blue
631        }
632
633        let device = self.tensor.device();
634        let tensor = Tensor::<B, 3>::from_data(TensorData::new(out, [3, h, w]), &device);
635        Ok(Image::new(tensor))
636    }
637
638    /// Converts an RGB image to HSL (Hue, Saturation, Lightness) color space.
639    /// Input must be a 3-channel image with values in [0, 1].
640    /// H is normalized to [0, 1] (from 360 degrees), S and L are in [0, 1].
641    pub fn rgb_to_hsl(&self) -> Result<Self> {
642        let dims = self.tensor.dims();
643        if dims[0] != 3 {
644            return Err(IrisError::InvalidParameter(
645                "Input must be a 3-channel RGB image".into(),
646            ));
647        }
648        let h = dims[1];
649        let w = dims[2];
650        let pixels = h * w;
651
652        let data = self.tensor.clone().into_data();
653        let flat: Vec<f32> = data.iter::<f32>().collect();
654        let mut out = vec![0.0f32; 3 * pixels];
655
656        for i in 0..pixels {
657            let r = flat[i] as f64;
658            let g = flat[pixels + i] as f64;
659            let b = flat[2 * pixels + i] as f64;
660
661            let max = r.max(g).max(b);
662            let min = r.min(g).min(b);
663            let l = (max + min) / 2.0;
664            let delta = max - min;
665
666            // Saturation
667            let s = if delta.abs() < 1e-10 {
668                0.0
669            } else if l < 0.5 {
670                delta / (max + min)
671            } else {
672                delta / (2.0 - max - min)
673            };
674
675            // Hue
676            let hue_deg = if delta.abs() < 1e-10 {
677                0.0
678            } else if (max - r).abs() < 1e-10 {
679                60.0 * (((g - b) / delta) % 6.0)
680            } else if (max - g).abs() < 1e-10 {
681                60.0 * (((b - r) / delta) + 2.0)
682            } else {
683                60.0 * (((r - g) / delta) + 4.0)
684            };
685
686            let hue_norm = if hue_deg < 0.0 {
687                (hue_deg + 360.0) / 360.0
688            } else {
689                hue_deg / 360.0
690            };
691
692            out[i] = hue_norm as f32;
693            out[pixels + i] = s.clamp(0.0, 1.0) as f32;
694            out[2 * pixels + i] = l.clamp(0.0, 1.0) as f32;
695        }
696
697        let device = self.tensor.device();
698        let tensor = Tensor::<B, 3>::from_data(TensorData::new(out, [3, h, w]), &device);
699        Ok(Image::new(tensor))
700    }
701
702    /// Converts an HSL image to RGB color space.
703    /// Input must be a 3-channel image where H is in [0, 1] (from 360 degrees),
704    /// S and L are in [0, 1].
705    /// Returns a 3-channel RGB image with values in [0, 1].
706    pub fn hsl_to_rgb(&self) -> Result<Self> {
707        let dims = self.tensor.dims();
708        if dims[0] != 3 {
709            return Err(IrisError::InvalidParameter(
710                "Input must be a 3-channel HSL image".into(),
711            ));
712        }
713        let h_dim = dims[1];
714        let w = dims[2];
715        let pixels = h_dim * w;
716
717        let data = self.tensor.clone().into_data();
718        let flat: Vec<f32> = data.iter::<f32>().collect();
719        let mut out = vec![0.0f32; 3 * pixels];
720
721        for i in 0..pixels {
722            let hue_deg = flat[i] as f64 * 360.0;
723            let s = flat[pixels + i] as f64;
724            let l = flat[2 * pixels + i] as f64;
725
726            let c = (1.0 - (2.0 * l - 1.0).abs()) * s;
727            let x = c * (1.0 - ((hue_deg / 60.0) % 2.0 - 1.0).abs());
728            let m = l - c / 2.0;
729
730            let (r, g, b) = if hue_deg < 60.0 {
731                (c, x, 0.0)
732            } else if hue_deg < 120.0 {
733                (x, c, 0.0)
734            } else if hue_deg < 180.0 {
735                (0.0, c, x)
736            } else if hue_deg < 240.0 {
737                (0.0, x, c)
738            } else if hue_deg < 300.0 {
739                (x, 0.0, c)
740            } else {
741                (c, 0.0, x)
742            };
743
744            out[i] = (r + m).clamp(0.0, 1.0) as f32;
745            out[pixels + i] = (g + m).clamp(0.0, 1.0) as f32;
746            out[2 * pixels + i] = (b + m).clamp(0.0, 1.0) as f32;
747        }
748
749        let device = self.tensor.device();
750        let tensor = Tensor::<B, 3>::from_data(TensorData::new(out, [3, h_dim, w]), &device);
751        Ok(Image::new(tensor))
752    }
753
754    /// Converts YCrCb to RGB color space.
755    pub fn ycrcb_to_rgb(&self) -> Result<Self> {
756        let dims = self.tensor.dims();
757        if dims[0] != 3 {
758            return Err(IrisError::InvalidParameter(
759                "Input must be 3-channel YCrCb".into(),
760            ));
761        }
762        let h = dims[1];
763        let w = dims[2];
764        let data = self.tensor.clone().into_data();
765        let flat: Vec<f32> = data.iter::<f32>().collect();
766        let pixels = h * w;
767        let mut out = vec![0.0f32; 3 * pixels];
768
769        for i in 0..pixels {
770            let y = flat[i] as f64;
771            let cr = flat[pixels + i] as f64 - 0.5;
772            let cb = flat[2 * pixels + i] as f64 - 0.5;
773
774            let r = y + 1.402 * cr;
775            let g = y - 0.714 * cr - 0.344 * cb;
776            let b = y + 1.772 * cb;
777
778            out[i] = r.clamp(0.0, 1.0) as f32;
779            out[pixels + i] = g.clamp(0.0, 1.0) as f32;
780            out[2 * pixels + i] = b.clamp(0.0, 1.0) as f32;
781        }
782
783        Ok(Image::new(Tensor::<B, 3>::from_data(
784            TensorData::new(out, [3, h, w]),
785            &self.tensor.device(),
786        )))
787    }
788}
789
790// Helper functions for color space conversions
791
792fn linearize(srgb: f32) -> f32 {
793    let v = srgb as f64;
794    if v <= 0.040_45 {
795        (v / 12.92) as f32
796    } else {
797        ((v + 0.055) / 1.055).powf(2.4) as f32
798    }
799}
800
801fn delinearize(lin: f32) -> f32 {
802    let v = lin as f64;
803    if v <= 0.003_130_8 {
804        (12.92 * v).clamp(0.0, 1.0) as f32
805    } else {
806        (1.055 * v.powf(1.0 / 2.4) - 0.055).clamp(0.0, 1.0) as f32
807    }
808}
809
810fn lab_f(t: f64) -> f64 {
811    let eps = 216.0 / 24_389.0;
812    let kappa = 24_389.0 / 27.0;
813    if t > eps {
814        t.cbrt()
815    } else {
816        (kappa * t + 16.0) / 116.0
817    }
818}
819
820fn lab_f_inv(t: f64) -> f64 {
821    let eps = 216.0 / 24_389.0;
822    let kappa = 24_389.0 / 27.0;
823    let t3 = t * t * t;
824    if t3 > eps {
825        t3
826    } else {
827        (116.0 * t - 16.0) / kappa
828    }
829}
830
831#[cfg(test)]
832mod tests {
833    use super::*;
834    use crate::test_helpers::{TestBackend, test_device};
835    use burn::tensor::TensorData;
836
837    #[test]
838    fn test_hsv_roundtrip() {
839        let device = test_device();
840        let flat_data = vec![
841            1.0, 0.0, 0.0, // Red
842            0.0, 1.0, 0.0, // Green
843            0.0, 0.0, 1.0, // Blue
844            1.0, 1.0, 0.0, // Yellow
845        ];
846        let tensor =
847            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 2, 2]), &device);
848        let rgb = Image::new(tensor);
849
850        let hsv = rgb.rgb_to_hsv().unwrap();
851        assert_eq!(hsv.shape(), [3, 2, 2]);
852
853        let back_rgb = hsv.hsv_to_rgb().unwrap();
854        assert_eq!(back_rgb.shape(), [3, 2, 2]);
855    }
856
857    #[test]
858    fn test_hls_roundtrip() {
859        let device = test_device();
860        let flat_data = vec![0.5f32; 3 * 4 * 4];
861        let tensor =
862            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 4, 4]), &device);
863        let rgb = Image::new(tensor);
864
865        let hls = rgb.rgb_to_hls().unwrap();
866        assert_eq!(hls.shape(), [3, 4, 4]);
867
868        let back_rgb = hls.hls_to_rgb().unwrap();
869        assert_eq!(back_rgb.shape(), [3, 4, 4]);
870    }
871
872    #[test]
873    fn test_split_merge() {
874        let device = test_device();
875        let flat_data = vec![0.3, 0.6, 0.9, 0.1, 0.4, 0.7];
876        let tensor =
877            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 1, 2]), &device);
878        let img = Image::new(tensor);
879
880        let channels = img.split_channels().unwrap();
881        assert_eq!(channels.len(), 3);
882
883        let merged = Image::merge_channels(&channels).unwrap();
884        assert_eq!(merged.shape(), [3, 1, 2]);
885    }
886
887    #[test]
888    fn test_xyz_roundtrip() {
889        let device = test_device();
890        let data = vec![0.5f32; 3 * 4 * 4];
891        let tensor = Tensor::<TestBackend, 3>::from_data(TensorData::new(data, [3, 4, 4]), &device);
892        let rgb = Image::new(tensor);
893        let xyz = rgb.rgb_to_xyz().unwrap();
894        assert_eq!(xyz.shape(), [3, 4, 4]);
895        let back = xyz.xyz_to_rgb().unwrap();
896        assert_eq!(back.shape(), [3, 4, 4]);
897    }
898
899    #[test]
900    fn test_lab_roundtrip() {
901        let device = test_device();
902        let data = vec![0.5f32; 3 * 4 * 4];
903        let tensor = Tensor::<TestBackend, 3>::from_data(TensorData::new(data, [3, 4, 4]), &device);
904        let rgb = Image::new(tensor);
905        let lab = rgb.rgb_to_lab().unwrap();
906        assert_eq!(lab.shape(), [3, 4, 4]);
907        let back = lab.lab_to_rgb().unwrap();
908        assert_eq!(back.shape(), [3, 4, 4]);
909    }
910
911    #[test]
912    fn test_yuv_roundtrip() {
913        let device = test_device();
914        let data = vec![0.5f32; 3 * 4 * 4];
915        let tensor = Tensor::<TestBackend, 3>::from_data(TensorData::new(data, [3, 4, 4]), &device);
916        let rgb = Image::new(tensor);
917        let yuv = rgb.rgb_to_yuv().unwrap();
918        assert_eq!(yuv.shape(), [3, 4, 4]);
919        let back = yuv.yuv_to_rgb().unwrap();
920        assert_eq!(back.shape(), [3, 4, 4]);
921    }
922
923    #[test]
924    fn test_ycrcb_roundtrip() {
925        let device = test_device();
926        let data = vec![0.5f32; 3 * 4 * 4];
927        let tensor = Tensor::<TestBackend, 3>::from_data(TensorData::new(data, [3, 4, 4]), &device);
928        let rgb = Image::new(tensor);
929        let ycrcb = rgb.rgb_to_ycrcb().unwrap();
930        assert_eq!(ycrcb.shape(), [3, 4, 4]);
931        let back = ycrcb.ycrcb_to_rgb().unwrap();
932        assert_eq!(back.shape(), [3, 4, 4]);
933    }
934
935    #[test]
936    fn test_cmyk_roundtrip() {
937        let device = test_device();
938        let flat_data = vec![
939            // R channel (4 pixels)
940            1.0, 0.0, 0.0, 0.5, // G channel (4 pixels)
941            0.0, 1.0, 0.0, 0.5, // B channel (4 pixels)
942            0.0, 0.0, 1.0, 0.5,
943        ];
944        let tensor =
945            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 2, 2]), &device);
946        let rgb = Image::new(tensor);
947
948        let cmyk = rgb.rgb_to_cmyk().unwrap();
949        assert_eq!(cmyk.shape(), [4, 2, 2]);
950
951        let back_rgb = cmyk.cmyk_to_rgb().unwrap();
952        assert_eq!(back_rgb.shape(), [3, 2, 2]);
953
954        // Verify roundtrip values
955        let orig_data = rgb.tensor.into_data();
956        let back_data = back_rgb.tensor.into_data();
957        let orig_vals: Vec<f32> = orig_data.iter::<f32>().collect();
958        let back_vals: Vec<f32> = back_data.iter::<f32>().collect();
959        for (a, b) in orig_vals.iter().zip(back_vals.iter()) {
960            assert!(
961                (a - b).abs() < 1e-5,
962                "CMYK roundtrip mismatch: {} vs {}",
963                a,
964                b
965            );
966        }
967    }
968
969    #[test]
970    fn test_hsl_roundtrip() {
971        let device = test_device();
972        let flat_data = vec![
973            1.0, 0.0, 0.0, // Red
974            0.0, 1.0, 0.0, // Green
975            0.0, 0.0, 1.0, // Blue
976            0.5, 0.5, 0.5, // Gray
977        ];
978        let tensor =
979            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 2, 2]), &device);
980        let rgb = Image::new(tensor);
981
982        let hsl = rgb.rgb_to_hsl().unwrap();
983        assert_eq!(hsl.shape(), [3, 2, 2]);
984
985        let back_rgb = hsl.hsl_to_rgb().unwrap();
986        assert_eq!(back_rgb.shape(), [3, 2, 2]);
987
988        // Verify roundtrip values
989        let orig_data = rgb.tensor.into_data();
990        let back_data = back_rgb.tensor.into_data();
991        let orig_vals: Vec<f32> = orig_data.iter::<f32>().collect();
992        let back_vals: Vec<f32> = back_data.iter::<f32>().collect();
993        for (a, b) in orig_vals.iter().zip(back_vals.iter()) {
994            assert!(
995                (a - b).abs() < 1e-5,
996                "HSL roundtrip mismatch: {} vs {}",
997                a,
998                b
999            );
1000        }
1001    }
1002
1003    #[test]
1004    fn test_color_invalid_channel() {
1005        let device = test_device();
1006        let data = vec![0.5f32; 4 * 4 * 4]; // 4 channels
1007        let tensor = Tensor::<TestBackend, 3>::from_data(TensorData::new(data, [4, 4, 4]), &device);
1008        let img = Image::new(tensor);
1009        assert!(img.rgb_to_hsv().is_err());
1010        assert!(img.rgb_to_xyz().is_err());
1011        assert!(img.rgb_to_cmyk().is_err());
1012        assert!(img.rgb_to_hsl().is_err());
1013
1014        // cmyk_to_rgb requires 4 channels, so 3-channel should fail
1015        let data3 = vec![0.5f32; 3 * 4 * 4];
1016        let tensor3 =
1017            Tensor::<TestBackend, 3>::from_data(TensorData::new(data3, [3, 4, 4]), &device);
1018        let img3 = Image::new(tensor3);
1019        assert!(img3.cmyk_to_rgb().is_err());
1020
1021        // hsl_to_rgb requires 3 channels, so 4-channel should fail
1022        assert!(img.hsl_to_rgb().is_err());
1023    }
1024}
iris/color/mod.rs

iris/color/
mod.rs
