nx_vitpose_preprocess.nx source

1// nx_vitpose_preprocess.nx -- image -> ViTPose input tensor: bilinear-resize an RGB image [sh,sw,3] (u8 0-255, 2// interleaved) to dh x dw and ImageNet-normalize into out [1,3,dh,dw] f32 (NCHW): out = (pix/255 - mean_c)/std_c. 3// The resize uses arbitrary-ratio bilinear via INTEGER coordinate math (src coord sy=(oy+0.5)*sh/dh-0.5 = rational 4// with denom 2*dh -> exact int floor + weights, align_corners=False, PyTorch-matching). Reusable for ANY real-image 5// inference (semantic validation of the ViTPose port + running it as a teacher for build-2). license_tier: ORIGINAL 6import "nx_syscalls.nx" 7import "nx_f32.nx" 8import "nx_f32_cvt.nx" 9import "nx_f32_div.nx" 10 11// rgb[(y*sw+x)*3+c] in 0..255 (int). mean/std are f32[3]. out is [1,3,dh,dw] f32 (NCHW). 12func vitpose_resize_normalize(rgb: *i64, sh: i64, sw: i64, dh: i64, dw: i64, mean: *i64, std: *i64, out: *i64) -> i64 { 13 let f255: i64 = nx_i32_to_f32(255) 14 let DENy: i64 = 2 * dh 15 let DENx: i64 = 2 * dw 16 let DEN2f: i64 = nx_i32_to_f32(DENy * DENx) 17 let plane: i64 = dh * dw 18 var oy: i64 = 0 19 while oy < dh { 20 let SY: i64 = (2 * oy + 1) * sh - dh 21 var qy: i64 = SY / DENy 22 var ry: i64 = SY - qy * DENy 23 if ry < 0 { qy = qy - 1; ry = ry + DENy } 24 var y0: i64 = qy; var y1: i64 = qy + 1 25 if y0 < 0 { y0 = 0 } 26 if y0 > sh - 1 { y0 = sh - 1 } 27 if y1 < 0 { y1 = 0 } 28 if y1 > sh - 1 { y1 = sh - 1 } 29 let iwy: i64 = DENy - ry 30 var ox: i64 = 0 31 while ox < dw { 32 let SX: i64 = (2 * ox + 1) * sw - dw 33 var qx: i64 = SX / DENx 34 var rx: i64 = SX - qx * DENx 35 if rx < 0 { qx = qx - 1; rx = rx + DENx } 36 var x0: i64 = qx; var x1: i64 = qx + 1 37 if x0 < 0 { x0 = 0 } 38 if x0 > sw - 1 { x0 = sw - 1 } 39 if x1 < 0 { x1 = 0 } 40 if x1 > sw - 1 { x1 = sw - 1 } 41 let iwx: i64 = DENx - rx 42 let w00: i64 = iwy * iwx; let w01: i64 = iwy * rx; let w10: i64 = ry * iwx; let w11: i64 = ry * rx 43 var c: i64 = 0 44 while c < 3 { 45 let v00: i64 = rgb[(y0*sw + x0)*3 + c] 46 let v01: i64 = rgb[(y0*sw + x1)*3 + c] 47 let v10: i64 = rgb[(y1*sw + x0)*3 + c] 48 let v11: i64 = rgb[(y1*sw + x1)*3 + c] 49 let num: i64 = w00*v00 + w01*v01 + w10*v10 + w11*v11 // integer weighted sum 50 let pixf: i64 = nx_f32_div(nx_i32_to_f32(num), DEN2f) // interpolated pixel value 0..255 51 let scaled: i64 = nx_f32_div(pixf, f255) // /255 52 out[c*plane + oy*dw + ox] = nx_f32_div(nx_f32_sub(scaled, mean[c]), std[c]) 53 c = c + 1 54 } 55 ox = ox + 1 56 } 57 oy = oy + 1 58 } 59 return 0 60}