nx_pose_preprocess.nx source
↩ module page · 183 lines · 9690 B
1// nx_pose_preprocess.nx -- ViTPose top-down preprocessing: real image + GT bbox -> vitpose_input.bin.
2// Pipeline (faithful to mmpose TopDownGetBboxCenterScale + TopDownAffine):
3// decode baseline/progressive JPEG (nx_jpeg_decode_rgb) -> packed RGB
4// -> aspect-fix bbox to 192:256 (3:4), pad x1.25 (mmpose default) -> crop box [x0,y0,cw,ch]
5// -> bilinear resize crop -> 192x256 -> ImageNet normalize ((v/255-mean)/std, RGB)
6// -> write CHW f32 LE (3*256*192*4 B) exactly as nx_vitpose_forward reads it.
7// Emits data/mp_crop.txt "x0 y0 cw ch" so postprocess maps heatmap(48x64)->image coords with the SAME box.
8// Config: data/mp_img_dest.txt (jpg path), data/mp_bbox.txt ("x y w h"). 100% sovereign decode+math.
9// verdict: NOT_YET_EVALUATED lineage_id: nx_pose_preprocess_v1
10import "nx_syscalls.nx"
11import "nx_itoa_lib.nx" // shared MSB-first emitter (zero-alloc)
12import "nx_jpeg_ascii.nx"
13import "nx_f32.nx"
14import "nx_f32_cvt.nx"
15import "nx_f32_div.nx"
16const K_MAGIC_4096: i64 = 4096
17const K_MAGIC_1024: i64 = 1024
18
19func pp_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return n }
20// MIGRATED to the shared emitter (debt 1785563586). The old body mmapped a scratch buffer
21// per call and never freed it. At PAGE granularity that is 4096B leaked PER CALL -- the
22// defect that took 28.5GB of a 36GB host in nx_ts_lumadiff (2MB input, ~3.66M calls).
23// nxi_* is MSB-first, allocates NOTHING, and emits identical bytes including the sign.
24func pp_wn(v: i64) -> i64 { nxi_out(v); return 0 }
25func pp_read(path: *u8, buf: *u8, cap: i64) -> i64 {
26 let fd: i64 = sys_openat_rd(path); if fd < 0 { return 0 - 1 }
27 var off: i64 = 0; var go: i64 = 1
28 while go == 1 { let n: i64 = sys_read(fd, ((buf as i64)+off) as *u8, cap-off); if n <= 0 { go = 0 } else { off = off+n; if off >= cap { go = 0 } } }
29 sys_close(fd); return off
30}
31func pp_cfg(path: *u8, out: *u8, cap: i64) -> i64 {
32 let fd: i64 = sys_openat_rd(path); if fd < 0 { out[0] = 0 as u8; return 0 - 1 }
33 let tmp: *u8 = sys_mmap(K_MAGIC_4096); let n: i64 = sys_read(fd, tmp, K_MAGIC_4096); sys_close(fd)
34 var k: i64 = 0; var go: i64 = 1
35 while go == 1 { if k >= n { go = 0 } else { let c: i64 = tmp[k] as i64
36 if c == 10 { go = 0 } else { if c == 13 { go = 0 } else { out[k] = tmp[k]; k = k+1; if k >= cap-1 { go = 0 } } } } }
37 out[k] = 0 as u8; return k
38}
39// parse up to 4 space-separated ints from a NUL-term string into v[4]; returns count.
40func pp_ints(s: *u8, v: *i64) -> i64 {
41 var i: i64 = 0; var c: i64 = 0
42 while c < 4 {
43 while s[i] == (32 as u8) { i = i + 1 }
44 if s[i] == 0 as u8 { c = 4 } else {
45 var val: i64 = 0; var any: i64 = 0
46 while s[i] >= (48 as u8) { if s[i] <= (57 as u8) { val = val*10 + ((s[i] as i64)-48); i = i+1; any = 1 } else { break } }
47 if any == 1 { v[c] = val; c = c + 1 } else { c = 4 }
48 }
49 }
50 return c
51}
52// truncate a NON-NEGATIVE f32 (bit-pattern in low32) to int; negative -> -1.
53func pp_f2i(x: i64) -> i64 {
54 if nx_f32_is_zero(x) == 1 { return 0 }
55 if ((x >> 31) & 1) == 1 { return 0 - 1 }
56 let e: i64 = ((x >> 23) & 0xFF) - 127
57 if e < 0 { return 0 }
58 let mant: i64 = (x & 0x7FFFFF) | 0x800000
59 if e >= 23 { return mant << (e - 23) }
60 return mant >> (23 - e)
61}
62
63func main() -> i64 {
64 let jpath: *u8 = sys_mmap(K_MAGIC_1024)
65 let bboxs: *u8 = sys_mmap(256)
66 pp_cfg("data/mp_img_dest.txt" as *u8, jpath, K_MAGIC_1024)
67 pp_cfg("data/mp_bbox.txt" as *u8, bboxs, 256)
68 let bv: *i64 = sys_mmap(8*4)
69 if pp_ints(bboxs, bv) < 4 { pp_w("ERR bbox\n" as *u8); sys_exit(1) }
70 let bx: i64 = bv[0]; let by: i64 = bv[1]; let bw: i64 = bv[2]; let bh: i64 = bv[3]
71 pp_w("jpg=" as *u8); pp_w(jpath); pp_w(" bbox=" as *u8); pp_wn(bx); pp_w("," as *u8); pp_wn(by); pp_w("," as *u8); pp_wn(bw); pp_w("," as *u8); pp_wn(bh); pp_w("\n" as *u8)
72
73 // ---- decode JPEG ----
74 let jb: *u8 = sys_mmap(4*K_MAGIC_1024*K_MAGIC_1024)
75 let jn: i64 = pp_read(jpath, jb, 4*K_MAGIC_1024*K_MAGIC_1024)
76 pp_w("jpg_bytes=" as *u8); pp_wn(jn); pp_w("\n" as *u8)
77 let rgbp: *i64 = sys_mmap(8); let wp: *i64 = sys_mmap(8); let hp: *i64 = sys_mmap(8)
78 let rc: i64 = nx_jpeg_decode_rgb(jb, jn, rgbp, wp, hp)
79 if rc != NX_JPEG_ASCII_OK { pp_w("ERR decode rc=" as *u8); pp_wn(rc); pp_w("\n" as *u8); sys_exit(1) }
80 let rgb: *u8 = rgbp[0] as *u8
81 let W: i64 = wp[0]; let H: i64 = hp[0]
82 pp_w("DECODED W=" as *u8); pp_wn(W); pp_w(" H=" as *u8); pp_wn(H); pp_w("\n" as *u8)
83
84 // ---- aspect-fix to 3:4 + pad x1.25 (mmpose) -> crop box ----
85 let cx: i64 = bx + bw/2; let cy: i64 = by + bh/2
86 var cw: i64 = 0; var ch: i64 = 0
87 if bw*256 > bh*192 { cw = bw; ch = (bw*256)/192 } else { ch = bh; cw = (bh*192)/256 }
88 cw = (cw*5)/4; ch = (ch*5)/4 // mmpose padding=1.25
89 let x0: i64 = cx - cw/2; let y0: i64 = cy - ch/2
90 pp_w("crop x0=" as *u8); pp_wn(x0); pp_w(" y0=" as *u8); pp_wn(y0); pp_w(" cw=" as *u8); pp_wn(cw); pp_w(" ch=" as *u8); pp_wn(ch); pp_w("\n" as *u8)
91
92 // f32 constants
93 let one: i64 = nx_i32_to_f32(1)
94 let half: i64 = nx_f32_div(one, nx_i32_to_f32(2))
95 let c255: i64 = nx_i32_to_f32(255)
96 let mR: i64 = nx_f32_div(nx_i32_to_f32(485), nx_i32_to_f32(1000))
97 let mG: i64 = nx_f32_div(nx_i32_to_f32(456), nx_i32_to_f32(1000))
98 let mB: i64 = nx_f32_div(nx_i32_to_f32(406), nx_i32_to_f32(1000))
99 let sR: i64 = nx_f32_div(nx_i32_to_f32(229), nx_i32_to_f32(1000))
100 let sG: i64 = nx_f32_div(nx_i32_to_f32(224), nx_i32_to_f32(1000))
101 let sB: i64 = nx_f32_div(nx_i32_to_f32(225), nx_i32_to_f32(1000))
102 let cwf: i64 = nx_i32_to_f32(cw); let chf: i64 = nx_i32_to_f32(ch)
103 let ratx: i64 = nx_f32_div(cwf, nx_i32_to_f32(192))
104 let raty: i64 = nx_f32_div(chf, nx_i32_to_f32(256))
105 let x0f: i64 = nx_i32_to_f32(x0); let y0f: i64 = nx_i32_to_f32(y0)
106
107 let iN: i64 = 3*256*192
108 let outb: *u8 = sys_mmap(iN*4 + 16)
109
110 var oy: i64 = 0
111 while oy < 256 {
112 let oyf: i64 = nx_f32_add(nx_i32_to_f32(oy), half)
113 let syf: i64 = nx_f32_add(nx_f32_neg(half), nx_f32_add(y0f, nx_f32_mul(oyf, raty)))
114 let iy: i64 = pp_f2i(syf)
115 let fy: i64 = nx_f32_add(syf, nx_f32_neg(nx_i32_to_f32(iy)))
116 let omfy: i64 = nx_f32_add(one, nx_f32_neg(fy))
117 var iy0: i64 = iy; if iy0 < 0 { iy0 = 0 } if iy0 > H-1 { iy0 = H-1 }
118 var iy1: i64 = iy+1; if iy1 < 0 { iy1 = 0 } if iy1 > H-1 { iy1 = H-1 }
119 var ox: i64 = 0
120 while ox < 192 {
121 let oxf: i64 = nx_f32_add(nx_i32_to_f32(ox), half)
122 let sxf: i64 = nx_f32_add(nx_f32_neg(half), nx_f32_add(x0f, nx_f32_mul(oxf, ratx)))
123 let ix: i64 = pp_f2i(sxf)
124 let fx: i64 = nx_f32_add(sxf, nx_f32_neg(nx_i32_to_f32(ix)))
125 let omfx: i64 = nx_f32_add(one, nx_f32_neg(fx))
126 var ix0: i64 = ix; if ix0 < 0 { ix0 = 0 } if ix0 > W-1 { ix0 = W-1 }
127 var ix1: i64 = ix+1; if ix1 < 0 { ix1 = 0 } if ix1 > W-1 { ix1 = W-1 }
128 let p00: i64 = (iy0*W+ix0)*3; let p10: i64 = (iy0*W+ix1)*3
129 let p01: i64 = (iy1*W+ix0)*3; let p11: i64 = (iy1*W+ix1)*3
130 var c: i64 = 0
131 while c < 3 {
132 let v00: i64 = nx_i32_to_f32(rgb[p00+c] as i64)
133 let v10: i64 = nx_i32_to_f32(rgb[p10+c] as i64)
134 let v01: i64 = nx_i32_to_f32(rgb[p01+c] as i64)
135 let v11: i64 = nx_i32_to_f32(rgb[p11+c] as i64)
136 let top: i64 = nx_f32_add(nx_f32_mul(v00, omfx), nx_f32_mul(v10, fx))
137 let bot: i64 = nx_f32_add(nx_f32_mul(v01, omfx), nx_f32_mul(v11, fx))
138 let v: i64 = nx_f32_add(nx_f32_mul(top, omfy), nx_f32_mul(bot, fy))
139 var mean: i64 = mR; var std: i64 = sR
140 if c == 1 { mean = mG; std = sG }
141 if c == 2 { mean = mB; std = sB }
142 let vn: i64 = nx_f32_div(nx_f32_add(nx_f32_div(v, c255), nx_f32_neg(mean)), std)
143 let idx: i64 = c*256*192 + oy*192 + ox
144 outb[idx*4+0] = (vn & 0xff) as u8
145 outb[idx*4+1] = ((vn >> 8) & 0xff) as u8
146 outb[idx*4+2] = ((vn >> 16) & 0xff) as u8
147 outb[idx*4+3] = ((vn >> 24) & 0xff) as u8
148 c = c + 1
149 }
150 ox = ox + 1
151 }
152 oy = oy + 1
153 }
154
155 // write vitpose_input.bin
156 let ofd: i64 = sys_openat_wr("/home/elderwesto/vitpose_input.bin" as *u8, 0x1a4)
157 if ofd < 0 { pp_w("ERR open out\n" as *u8); sys_exit(1) }
158 var woff: i64 = 0
159 while woff < iN*4 { let ww: i64 = sys_write(ofd, ((outb as i64)+woff) as *u8, iN*4-woff); if ww <= 0 { woff = iN*4 } else { woff = woff+ww } }
160 sys_close(ofd)
161 pp_w("wrote vitpose_input.bin bytes=" as *u8); pp_wn(iN*4); pp_w("\n" as *u8)
162
163 // emit crop box for postprocess
164 let cbuf: *u8 = sys_mmap(128); var ck: i64 = 0
165 // format "x0 y0 cw ch\n"
166 let fnums: *i64 = sys_mmap(8*4); fnums[0]=x0; fnums[1]=y0; fnums[2]=cw; fnums[3]=ch
167 var fi: i64 = 0
168 while fi < 4 {
169 var m: i64 = fnums[fi]; let tb: *u8 = sys_mmap(28); var tk: i64 = 0
170 if m == 0 { tb[0]=48 as u8; tk=1 }
171 if m < 0 { cbuf[ck]=45 as u8; ck=ck+1; m=0-m }
172 while m > 0 { tb[tk]=(48+(m%10)) as u8; m=m/10; tk=tk+1 }
173 var ti: i64 = 0; while ti < tk { cbuf[ck]=tb[tk-1-ti]; ck=ck+1; ti=ti+1 }
174 cbuf[ck]=32 as u8; ck=ck+1; fi=fi+1
175 }
176 cbuf[ck]=10 as u8; ck=ck+1
177 let cfd: i64 = sys_openat_wr("data/mp_crop.txt" as *u8, 0x1a4)
178 if cfd >= 0 { sys_write(cfd, cbuf, ck); sys_close(cfd) }
179
180 pp_w("PREPROCESS OK -> input.bin ready + crop box emitted\n" as *u8)
181 sys_exit(0)
182 return 0
183}