diff --git a/src-tauri/src/ia/mod.rs b/src-tauri/src/ia/mod.rs index 13d8640..976c7a2 100644 --- a/src-tauri/src/ia/mod.rs +++ b/src-tauri/src/ia/mod.rs @@ -255,17 +255,17 @@ fn rellenar(rgba: &[u8], ancho: usize, alto: usize) -> Result, String> { let Some(a) = m.as_mut() else { return Err("La IA no está preparada.".into()); }; - let (mw, mh) = relleno::tamano_modelo(ancho, alto); + let p = relleno::plan(ancho, alto); let t0 = Instant::now(); - let r = a.motor.correr_forma(&[1, 4, mh, mw], relleno::entrada(rgba, ancho, alto, mw, mh)); + let r = a.motor.correr_forma(&[1, 4, p.ph, p.pw], relleno::entrada(rgba, ancho, alto, p)); a.usado = Instant::now(); match r { - Ok((_, planos)) if planos.len() == 3 * mw * mh => { - crate::salud::anotar("INFO", "ia", &format!("relleno de {ancho}×{alto} (a {mw}×{mh}) en {:.1} s", t0.elapsed().as_secs_f32())); - let rgb = relleno::salida(&planos, mw, mh, rgba, ancho, alto); + Ok((_, planos)) if planos.len() == 3 * p.pw * p.ph => { + crate::salud::anotar("INFO", "ia", &format!("relleno de {ancho}×{alto} (a {}×{} en {}×{}) en {:.1} s", p.sw, p.sh, p.pw, p.ph, t0.elapsed().as_secs_f32())); + let rgb = relleno::salida(&planos, p, rgba, ancho, alto); Ok(imagen::respuesta(ancho as u32, alto as u32, &rgb)) } - Ok((forma, _)) => Err(format!("El modelo contestó {forma:?} y se esperaba [1, 3, {mh}, {mw}].")), + Ok((forma, _)) => Err(format!("El modelo contestó {forma:?} y se esperaba [1, 3, {}, {}].", p.ph, p.pw)), Err(e) => { *m = None; crate::salud::anotar("ERROR", "ia", &format!("no se pudo rellenar: {e}")); @@ -357,12 +357,12 @@ mod pruebas { rgba[i + 3] = if hueco { 255 } else { 0 }; } } - let (mw, mh) = relleno::tamano_modelo(w, h); + let p = relleno::plan(w, h); let t = Instant::now(); - let (_, planos) = m.correr_forma(&[1, 4, mh, mw], relleno::entrada(&rgba, w, h, mw, mh)).unwrap(); - let rgb = relleno::salida(&planos, mw, mh, &rgba, w, h); + let (_, planos) = m.correr_forma(&[1, 4, p.ph, p.pw], relleno::entrada(&rgba, w, h, p)).unwrap(); + let rgb = relleno::salida(&planos, p, &rgba, w, h); let c = (h / 2 * w + w / 2) * 3; - eprintln!("{w}×{h} (a {mw}×{mh}): {:.2} s, centro {:?}, esquina {:?}", t.elapsed().as_secs_f32(), &rgb[c..c + 3], &rgb[..3]); + eprintln!("{w}×{h} ({p:?}): {:.2} s, centro {:?}, esquina {:?}", t.elapsed().as_secs_f32(), &rgb[c..c + 3], &rgb[..3]); // fuera del hueco, lo que había assert_eq!(&rgb[..3], &rgba[..3]); } diff --git a/src-tauri/src/ia/relleno.rs b/src-tauri/src/ia/relleno.rs index 218c153..c29cfa3 100644 --- a/src-tauri/src/ia/relleno.rs +++ b/src-tauri/src/ia/relleno.rs @@ -2,9 +2,14 @@ // // La página (lib/revelado/relleno y revelado-relleno de labstream-os) manda el trozo de la foto // alrededor de lo pintado, en 8 bits, con el hueco en el alfa (255 = rellenar). Aquí se lleva al -// tamaño del modelo —el del trozo, en múltiplos de 32, con 1024 px de lado como mucho: LaMa -// generaliza bien hasta ahí y en la CPU tarda 10 s; a 1536, 75—, se rellena, se vuelve a su tamaño y -// se devuelve el trozo con el hueco relleno y lo demás tal cual estaba (byte a byte). +// tamaño del modelo, se rellena, se vuelve a su tamaño y se devuelve el trozo con el hueco relleno y +// lo demás tal cual estaba (byte a byte). +// +// El tamaño (`plan`): LaMa hace FFT de lo que ve, y la de ONNX Runtime en la CPU solo es rápida con +// lados que son potencias de 2: 512 × 512 son 2,7 s y 704 × 448 —casi los mismos píxeles— 10 s +// (14-oct-2026). Así que el trozo se reduce lo justo para caber en 512 (o en 1024 si pasa de 768 px: +// LaMa generaliza bien hasta ahí, 10 s) y se completa hasta la potencia de 2 con su propio borde en +// espejo (sin hueco: es contexto). Lo de fuera de la foto reducida se descarta al volver. // // El modelo (LaMa, Apache-2.0; la exportación ONNX de tamaño libre de g-ronimo): una entrada // [1, 4, alto, ancho] —RGB 0..1 con el hueco a cero, y el hueco 0/1— y una salida [1, 3, alto, ancho] @@ -13,13 +18,23 @@ /// El lado mayor al que se rellena. pub const LADO_MAX: usize = 1024; -/// El tamaño al que se rellena un trozo de `ancho` × `alto`: el suyo (o reducido a LADO_MAX de lado -/// mayor), en múltiplos de 32 y con 64 como mínimo. -pub fn tamano_modelo(ancho: usize, alto: usize) -> (usize, usize) { - let mayor = ancho.max(alto) as f64; - let k = if mayor > LADO_MAX as f64 { LADO_MAX as f64 / mayor } else { 1.0 }; - let a32 = |v: f64| ((((v * k) / 32.0).round() as usize).max(2)) * 32; - (a32(ancho as f64), a32(alto as f64)) +/// Cómo entra un trozo de `ancho` × `alto` al modelo: reducido a `sw` × `sh` y completado en espejo +/// hasta `pw` × `ph` (potencias de 2). +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub struct Plan { + pub sw: usize, + pub sh: usize, + pub pw: usize, + pub ph: usize, +} + +pub fn plan(ancho: usize, alto: usize) -> Plan { + let mayor = ancho.max(alto); + let lado = if mayor <= 768 { 512 } else { LADO_MAX }; + let k = (lado as f64 / mayor as f64).min(1.0); + let sw = ((ancho as f64 * k).round() as usize).clamp(8, lado); + let sh = ((alto as f64 * k).round() as usize).clamp(8, lado); + Plan { sw, sh, pw: sw.next_power_of_two().max(64), ph: sh.next_power_of_two().max(64) } } /// Remuestrea `de` (`canales` por punto, `aw` × `ah`) a `bw` × `bh` entre sus cuatro vecinos. @@ -46,8 +61,18 @@ pub fn remuestrear(de: &[f32], aw: usize, ah: usize, canales: usize, bw: usize, out } -/// La entrada del modelo (planos [4, mh, mw]) a partir del RGBA de la página, y el hueco a su tamaño. -pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, mw: usize, mh: usize) -> Vec { +/// El índice en espejo de `i` dentro de 0..n (…2 1 0 1 2…n-2 n-1 n-2…). +fn espejo(i: usize, n: usize) -> usize { + if n <= 1 { + return 0; + } + let periodo = 2 * (n - 1); + let r = i % periodo; + if r < n { r } else { periodo - r } +} + +/// La entrada del modelo (planos [4, ph, pw]) a partir del RGBA de la página. +pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, p: Plan) -> Vec { let n = ancho * alto; let mut rgb = Vec::with_capacity(n * 3); let mut hueco = Vec::with_capacity(n); @@ -55,32 +80,42 @@ pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, mw: usize, mh: usize) -> rgb.extend_from_slice(&[rgba[i * 4] as f32 / 255.0, rgba[i * 4 + 1] as f32 / 255.0, rgba[i * 4 + 2] as f32 / 255.0]); hueco.push(if rgba[i * 4 + 3] >= 128 { 1.0 } else { 0.0 }); } - let rgb = remuestrear(&rgb, ancho, alto, 3, mw, mh); - let hueco = remuestrear(&hueco, ancho, alto, 1, mw, mh); - let m = mw * mh; + let rgb = remuestrear(&rgb, ancho, alto, 3, p.sw, p.sh); + let hueco = remuestrear(&hueco, ancho, alto, 1, p.sw, p.sh); + let m = p.pw * p.ph; let mut out = vec![0f32; 4 * m]; - for i in 0..m { - // lo que toca el hueco, aunque sea un poco al reducir, es hueco - let h = if hueco[i] > 0.02 { 1.0 } else { 0.0 }; - for c in 0..3 { - out[c * m + i] = if h > 0.0 { 0.0 } else { rgb[i * 3 + c] }; + for y in 0..p.ph { + let dentro_y = y < p.sh; + let sy = espejo(y, p.sh); + for x in 0..p.pw { + let dentro = dentro_y && x < p.sw; + let sx = espejo(x, p.sw); + let s = sy * p.sw + sx; + // lo que toca el hueco, aunque sea un poco al reducir, es hueco; el espejo, nunca + let h = if dentro && hueco[s] > 0.02 { 1.0 } else { 0.0 }; + let i = y * p.pw + x; + for c in 0..3 { + out[c * m + i] = if h > 0.0 { 0.0 } else { rgb[s * 3 + c] }; + } + out[3 * m + i] = h; } - out[3 * m + i] = h; } out } -/// La salida del modelo (planos [3, mh, mw], 0..1) de vuelta a `ancho` × `alto`: dentro del hueco lo +/// La salida del modelo (planos [3, ph, pw], 0..1) de vuelta a `ancho` × `alto`: dentro del hueco lo /// rellenado; fuera, el RGB que llegó, tal cual. RGB 8 bits. -pub fn salida(planos: &[f32], mw: usize, mh: usize, rgba: &[u8], ancho: usize, alto: usize) -> Vec { - let m = mw * mh; - let mut intercalado = vec![0f32; m * 3]; - for i in 0..m { - for c in 0..3 { - intercalado[i * 3 + c] = planos[c * m + i]; +pub fn salida(planos: &[f32], p: Plan, rgba: &[u8], ancho: usize, alto: usize) -> Vec { + let m = p.pw * p.ph; + let mut intercalado = vec![0f32; p.sw * p.sh * 3]; + for y in 0..p.sh { + for x in 0..p.sw { + for c in 0..3 { + intercalado[(y * p.sw + x) * 3 + c] = planos[c * m + y * p.pw + x]; + } } } - let grande = remuestrear(&intercalado, mw, mh, 3, ancho, alto); + let grande = remuestrear(&intercalado, p.sw, p.sh, 3, ancho, alto); let n = ancho * alto; let mut out = vec![0u8; n * 3]; for i in 0..n { @@ -97,26 +132,31 @@ mod pruebas { use super::*; #[test] - fn el_tamano_del_modelo() { - assert_eq!(tamano_modelo(500, 300), (512, 288)); - assert_eq!(tamano_modelo(2048, 1024), (1024, 512)); - assert_eq!(tamano_modelo(20, 10), (64, 64)); - let (w, h) = tamano_modelo(1600, 1067); - assert!(w <= 1024 && w % 32 == 0 && h % 32 == 0); + fn el_plan_reduce_lo_justo_y_completa_hasta_potencias_de_2() { + assert_eq!(plan(700, 460), Plan { sw: 512, sh: 336, pw: 512, ph: 512 }); + assert_eq!(plan(500, 300), Plan { sw: 500, sh: 300, pw: 512, ph: 512 }); + assert_eq!(plan(1400, 933), Plan { sw: 1024, sh: 682, pw: 1024, ph: 1024 }); + assert_eq!(plan(200, 100), Plan { sw: 200, sh: 100, pw: 256, ph: 128 }); + assert_eq!(plan(20, 10), Plan { sw: 20, sh: 10, pw: 64, ph: 64 }); + assert_eq!(espejo(5, 4), 1); + assert_eq!(espejo(3, 4), 3); } #[test] fn fuera_del_hueco_queda_lo_que_habia() { // 2 × 1: el primero es hueco let rgba = [10, 20, 30, 255, 40, 50, 60, 0]; - let e = entrada(&rgba, 2, 1, 64, 64); - assert_eq!(e.len(), 4 * 64 * 64); - // a la izquierda (el hueco) el RGB va a cero y el canal del hueco a 1 + let p = plan(2, 1); + let e = entrada(&rgba, 2, 1, p); + let m = p.pw * p.ph; + assert_eq!(e.len(), 4 * m); + // el hueco (0,0): RGB a cero y el canal del hueco a 1; el espejo (fuera de 2 × 1) no es hueco assert_eq!(e[0], 0.0); - assert_eq!(e[3 * 64 * 64], 1.0); + assert_eq!(e[3 * m], 1.0); + assert_eq!(e[3 * m + 10], 0.0); // lo de la red: todo gris - let planos = vec![0.5f32; 3 * 64 * 64]; - let s = salida(&planos, 64, 64, &rgba, 2, 1); + let planos = vec![0.5f32; 3 * m]; + let s = salida(&planos, p, &rgba, 2, 1); assert_eq!(&s[..3], &[128, 128, 128]); assert_eq!(&s[3..], &[40, 50, 60]); }