`nx_coco_oks.nx` source

1// nx_coco_oks.nx -- the REAL COCO industry metric OKS (Object Keypoint Similarity), replacing our self-validating
2// PCK-vs-our-own-teacher. OKS = Sum_i exp(-d_i^2 / (2 s^2 k_i^2)) * [v_i>0] / Sum_i [v_i>0], where d_i = pred-GT
3// distance, s^2 = the person's segment AREA (from the COCO annotation), k_i = 2*sigma_i, sigma_i the OFFICIAL COCO
4// per-keypoint constants. This is what AP@[.5:.95] thresholds -> the number the field grades pose by. Needs
5// GROUND-TRUTH annotations (area + GT keypoints) to be meaningful -- it is NOT computable against our own port.
6// license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_f32.nx"
9import "nx_f32_cvt.nx"
10import "nx_f32_div.nx"
11import "nx_f32_exp.nx"
12
13// the 17 official COCO keypoint sigmas x1000 (nose, eyes, ears, shoulders, elbows, wrists, hips, knees, ankles).
14func oks_sigma1000(i: i64) -> i64 {
15    let s: *i64 = sys_mmap(8*17) as *i64
16    s[0]=26; s[1]=25; s[2]=25; s[3]=35; s[4]=35; s[5]=79; s[6]=79; s[7]=72; s[8]=72
17    s[9]=62; s[10]=62; s[11]=107; s[12]=107; s[13]=87; s[14]=87; s[15]=89; s[16]=89
18    return s[i]
19}
20
21// px,py = predicted (pixel); gx,gy,vis = ground-truth keypoints + visibility; area = person segment area (f32).
22// Returns OKS in permille (0..1000).
23func coco_oks(px: *i64, py: *i64, gx: *i64, gy: *i64, vis: *i64, n: i64, area: i64) -> i64 {
24    let eight: i64 = nx_i32_to_f32(8)
25    var num: i64 = 0        // +0.0
26    var den: i64 = 0
27    var i: i64 = 0
28    while i < n {
29        if vis[i] > 0 {
30            den = den + 1
31            let dx: i64 = px[i] - gx[i]
32            let dy: i64 = py[i] - gy[i]
33            let d2: i64 = nx_i32_to_f32(dx*dx + dy*dy)
34            let sig: i64 = nx_f32_div(nx_i32_to_f32(oks_sigma1000(i)), nx_i32_to_f32(1000))   // sigma_i
35            let k2: i64 = nx_f32_mul(nx_f32_mul(nx_i32_to_f32(2), sig), nx_f32_mul(nx_i32_to_f32(2), sig))  // (2 sigma)^2
36            let denom: i64 = nx_f32_mul(nx_f32_mul(nx_i32_to_f32(2), area), k2)   // 2 s^2 k^2
37            let e: i64 = nx_f32_exp(nx_f32_neg(nx_f32_div(d2, denom)))
38            num = nx_f32_add(num, e)
39        }
40        i = i + 1
41    }
42    if den == 0 { return 0 }
43    let oks: i64 = nx_f32_div(num, nx_i32_to_f32(den))   // 0..1
44    // -> permille via *1000 then truncate
45    let scaled: i64 = nx_f32_mul(oks, nx_i32_to_f32(1000))
46    // truncate f32 -> int (non-negative, < 1001)
47    if nx_f32_is_zero(scaled) == 1 { return 0 }
48    let exp: i64 = ((scaled >> 23) & 0xFF) - 127
49    if exp < 0 { return 0 }
50    let mant: i64 = (scaled & 0x7FFFFF) | 0x800000
51    if exp >= 23 { return mant << (exp - 23) }
52    return mant >> (23 - exp)
53}
nx_coco_oks.nx source

`nx_coco_oks.nx` source
