File size: 6,140 Bytes
06e4b6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | //! `jepa::patches` — extraído do monólito na v0.10.0
//! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2).
use super::*;
// ---------------------------------------------------------------------------
// Parte 3 — adaptadores de modalidade
// ---------------------------------------------------------------------------
/// (C, H, W) → patches (H/ph · W/pw, C·ph·pw) — Image-JEPA.
pub fn em_patches_2d(img: &Tensor, ph: usize, pw: usize) -> Result<Tensor> {
let (c, h, w) = (img.dim(0)?, img.dim(1)?, img.dim(2)?);
let dev = img.device();
let mut linhas: Vec<Tensor> = Vec::new();
for y in (0..h).step_by(ph) {
for x in (0..w).step_by(pw) {
let hh = ph.min(h - y);
let ww = pw.min(w - x);
let patch = img.narrow(0, 0, c)?.narrow(1, y, hh)?.narrow(2, x, ww)?;
// padding para (C, ph, pw) se borda
let patch = if hh < ph || ww < pw {
let mut buf = patch.flatten_all()?.to_vec1::<f32>()?;
let falta = c * ph * pw - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
Tensor::from_vec(buf, (c * ph * pw,), &dev)?
} else {
patch.flatten_all()?
};
linhas.push(patch);
}
}
Ok(Tensor::stack(&linhas, 0)?)
}
/// (T, H, W) → tubos (T/pt · H/ph · W/pw, pt·ph·pw) — V-JEPA 2.
pub fn em_tubos_3d(vid: &Tensor, pt: usize, ph: usize, pw: usize) -> Result<Tensor> {
let (t, h, w) = (vid.dim(0)?, vid.dim(1)?, vid.dim(2)?);
let dev = vid.device();
let mut linhas: Vec<Tensor> = Vec::new();
for tt in (0..t).step_by(pt) {
let tl = pt.min(t - tt);
for y in (0..h).step_by(ph) {
for x in (0..w).step_by(pw) {
let hh = ph.min(h - y);
let ww = pw.min(w - x);
let tubo = vid
.narrow(0, tt, tl)?
.narrow(1, y, hh)?
.narrow(2, x, ww)?
.flatten_all()?;
let mut buf = tubo.to_vec1::<f32>()?;
let falta = pt * ph * pw - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
linhas.push(Tensor::from_vec(buf, (pt * ph * pw,), &dev)?);
}
}
}
Ok(Tensor::stack(&linhas, 0)?)
}
/// (F, τ) → patches (F/pf · τ/pt, pf·pt) — Audio-JEPA.
pub fn em_patches_espectrograma(spec: &Tensor, pf: usize, pt: usize) -> Result<Tensor> {
let (f, tau) = (spec.dim(0)?, spec.dim(1)?);
let dev = spec.device();
let mut linhas: Vec<Tensor> = Vec::new();
for i in (0..f).step_by(pf) {
for j in (0..tau).step_by(pt) {
let ff = pf.min(f - i);
let tt = pt.min(tau - j);
let p = spec.narrow(0, i, ff)?.narrow(1, j, tt)?.flatten_all()?;
let mut buf = p.to_vec1::<f32>()?;
let falta = pf * pt - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
linhas.push(Tensor::from_vec(buf, (pf * pt,), &dev)?);
}
}
Ok(Tensor::stack(&linhas, 0)?)
}
/// (C, T) → patches canal-independentes (C · T/pt, pt) — T-JEPA/Brain-JEPA.
pub fn em_patches_serie(ts: &Tensor, pt: usize) -> Result<Tensor> {
let (c, t) = (ts.dim(0)?, ts.dim(1)?);
let dev = ts.device();
let mut linhas: Vec<Tensor> = Vec::new();
for ch in 0..c {
for j in (0..t).step_by(pt) {
let tt = pt.min(t - j);
let p = ts.narrow(0, ch, 1)?.narrow(1, j, tt)?.flatten_all()?;
let mut buf = p.to_vec1::<f32>()?;
let falta = pt - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
linhas.push(Tensor::from_vec(buf, (pt,), &dev)?);
}
}
Ok(Tensor::stack(&linhas, 0)?)
}
/// (X, Y, Z) → voxels (X/px · Y/py · Z/pz, px·py·pz) — 3D-JEPA.
pub fn em_voxels(vol: &Tensor, px: usize, py: usize, pz: usize) -> Result<Tensor> {
let (x, y, z) = (vol.dim(0)?, vol.dim(1)?, vol.dim(2)?);
let dev = vol.device();
let mut linhas: Vec<Tensor> = Vec::new();
for i in (0..x).step_by(px) {
for j in (0..y).step_by(py) {
for k in (0..z).step_by(pz) {
let lx = px.min(x - i);
let ly = py.min(y - j);
let lz = pz.min(z - k);
let p = vol.narrow(0, i, lx)?.narrow(1, j, ly)?.narrow(2, k, lz)?.flatten_all()?;
let mut buf = p.to_vec1::<f32>()?;
let falta = px * py * pz - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
linhas.push(Tensor::from_vec(buf, (px * py * pz,), &dev)?);
}
}
}
Ok(Tensor::stack(&linhas, 0)?)
}
/// Nuvem de pontos (n, d) → grupos FPS (k, m·d): k centros por FPS, cada
/// grupo = os m vizinhos mais próximos do centro (concatenados) — Point-JEPA.
pub fn em_grupos_pontos(pc: &Tensor, k: usize, m: usize, rng: &mut RngDeterministico) -> Result<Tensor> {
let n = pc.dim(0)?;
let d = pc.dim(1)?;
let coords = pc.to_vec2::<f32>()?;
let coords_f: Vec<Vec<f64>> = coords.iter().map(|r| r.iter().map(|&v| v as f64).collect()).collect();
let centros = pontos_mais_distantes(&coords_f, k, rng);
let dev = pc.device();
let mut linhas: Vec<Tensor> = Vec::new();
for &ci in ¢ros {
// m vizinhos mais próximos do centro ci (incluindo ele próprio)
let mut ord: Vec<usize> = (0..n).collect();
ord.sort_by(|&a, &b| {
let da = distancia_eucl(&coords_f[a], &coords_f[ci]);
let db = distancia_eucl(&coords_f[b], &coords_f[ci]);
da.partial_cmp(&db).unwrap()
});
let mut buf: Vec<f32> = Vec::with_capacity(m * d);
for &idx in ord.iter().take(m) {
for j in 0..d {
buf.push(coords[idx][j]);
}
}
let falta = m * d - buf.len();
buf.extend(std::iter::repeat(0f32).take(falta));
linhas.push(Tensor::from_vec(buf, (m * d,), &dev)?);
}
Ok(Tensor::stack(&linhas, 0)?)
}
|