nx_vitpose_prep.nx source
↩ module page · 59 lines · 3764 B
1// nx_vitpose_prep.nx -- decode the fetched JPEG + preprocess into the ViTPose input tensor, written to a file the
2// forward reads. Reads the image path from data/mp_img_dest.txt: nx_jpeg_decode_rgb -> u8 RGB [h,w,3] -> i64 ->
3// nx_vitpose_preprocess(resize 256x192 + ImageNet normalize) -> [1,3,256,192] f32 -> models/vitpose_input.bin
4// (f32 as u32le). expect_exit: 0
5import "nx_syscalls.nx"
6import "nx_f32.nx"
7import "nx_f32_cvt.nx"
8import "nx_f32_div.nx"
9import "nx_jpeg_ascii.nx"
10import "nx_vitpose_preprocess.nx"
11const K_MAGIC_4096: i64 = 4096
12const K_MAGIC_4194304: i64 = 4194304
13
14func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
15func wn(v: i64) -> i64 { var m: i64=v; if m<0{w("-" as *u8);m=0-m} let t:*u8=sys_mmap(24); var k:i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i:i64=0; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1} sys_write(1,o,k); return 0 }
16func rdcfg(path: *u8, buf: *u8, cap: i64) -> i64 { let fd: i64=sys_openat_rd(path); if fd<0{return 0-1} var n: i64=sys_read(fd,buf,cap-1); sys_close(fd); if n<0{return 0-1} var go: i64=1; while go==1{go=0; if n>0{ let c: i64=buf[n-1] as i64; if c==10{n=n-1;go=1} if c==13{n=n-1;go=1} if c==32{n=n-1;go=1} }} buf[n]=0 as u8; return n }
17func mk(n: i64, d: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(n), nx_i32_to_f32(d)) }
18
19func main() -> i64 {
20 let path: *u8=sys_mmap(K_MAGIC_4096)
21 if rdcfg("data/mp_img_dest.txt" as *u8, path, K_MAGIC_4096)<0 { w("no img path\n" as *u8); sys_exit(1); return 1 }
22 let fd: i64=sys_openat_rd(path); if fd<0{ w("open jpeg FAILED\n" as *u8); sys_exit(1); return 1 }
23 let JCAP: i64=K_MAGIC_4194304; let jpeg: *u8=sys_mmap(JCAP); var jl: i64=0; var go: i64=1
24 while go==1 { let n: i64=sys_read(fd, ((jpeg as i64)+jl) as *u8, JCAP-jl); if n<=0{go=0}else{jl=jl+n} }
25 sys_close(fd)
26 w("jpeg bytes=" as *u8); wn(jl); w("\n" as *u8)
27
28 let orgb: *i64=sys_mmap(8) as *i64; let ow: *i64=sys_mmap(8) as *i64; let oh: *i64=sys_mmap(8) as *i64
29 let drc: i64=nx_jpeg_decode_rgb(jpeg, jl, orgb, ow, oh)
30 if drc != 0 { w("jpeg decode FAILED rc=" as *u8); wn(drc); w("\n" as *u8); sys_exit(1); return 1 }
31 let rgbu: *u8=(orgb[0]) as *u8; let W: i64=ow[0]; let H: i64=oh[0]
32 w("decoded W=" as *u8); wn(W); w(" H=" as *u8); wn(H); w("\n" as *u8)
33
34 // u8 RGB -> i64
35 let rgb: *i64=sys_mmap(8*W*H*3) as *i64
36 var i: i64=0; while i<W*H*3 { rgb[i]=(rgbu[i] as i64)&0xff; i=i+1 }
37
38 let mean: *i64=sys_mmap(8*3) as *i64; let std: *i64=sys_mmap(8*3) as *i64
39 mean[0]=mk(485,1000); mean[1]=mk(456,1000); mean[2]=mk(406,1000)
40 std[0]=mk(229,1000); std[1]=mk(224,1000); std[2]=mk(225,1000)
41 let inp: *i64=sys_mmap(8*3*256*192) as *i64
42 vitpose_resize_normalize(rgb, H, W, 256, 192, mean, std, inp)
43
44 // write [1,3,256,192] f32 as u32le -> models/vitpose_input.bin
45 let N: i64=3*256*192
46 let ob: *u8=sys_mmap(N*4+16)
47 i=0; while i<N { let v: i64=inp[i]; ob[i*4]=(v&0xff) as u8; ob[i*4+1]=((v>>8)&0xff) as u8; ob[i*4+2]=((v>>16)&0xff) as u8; ob[i*4+3]=((v>>24)&0xff) as u8; i=i+1 }
48 let ofd: i64=sys_openat_wr("/mnt/c/Users/elder/nishi-core/models/vitpose_input.bin" as *u8, 0x1a4)
49 if ofd<0 { w("open out FAILED\n" as *u8); sys_exit(1); return 1 }
50 var wr: i64=0; while wr<N*4 { let k: i64=sys_write(ofd, ((ob as i64)+wr) as *u8, N*4-wr); if k<=0{wr=N*4}else{wr=wr+k} }
51 sys_close(ofd)
52 w("wrote vitpose_input.bin (" as *u8); wn(N); w(" f32 values)\n" as *u8)
53 // sample a few normalized values
54 w(" input[0..3]= " as *u8)
55 i=0; while i<3 { let v: i64=inp[i]; let neg: i64=nx_f32_lt(v,0); if neg==1{w("-" as *u8)} wn(0); i=i+1 }
56 w(" (see forward for kpts)\n" as *u8)
57 sys_exit(0)
58 return 0
59}