nx_h264_slice_extract.nx source
↩ module page · 339 lines · 17486 B
1// nx_h264_slice_extract.nx -- locate + parse a REAL coded slice from realtest.mp4.
2// Chains the whole stack on real bytes: avcC -> SPS+PPS params; video stbl ->
3// stco[0]/stsz[0] -> sample 0 file offset+size; AVCC NAL-split that sample; find
4// the slice NAL; parse its slice header with the real SPS/PPS context. This is
5// the bridge from "front-end parses config" to "we are holding a real coded
6// macroblock slice" -- the input CAVLC will decode.
7// license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_mp4_demux.nx"
10import "nx_mp4_stbl.nx"
11import "nx_h264_bits.nx"
12import "nx_h264_sps.nx"
13import "nx_h264_pps.nx"
14import "nx_h264_nal.nx"
15import "nx_h264_slice.nx"
16import "nx_h264_mb.nx"
17import "nx_h264_mb_pred.nx"
18import "nx_h264_coeff_token.nx"
19import "nx_h264_cavlc_level.nx"
20import "nx_h264_residual.nx"
21import "nx_h264_zigzag.nx"
22import "nx_h264_hadamard.nx"
23import "nx_h264_idct.nx"
24import "nx_h264_intra16.nx"
25import "nx_h264_deblock.nx"
26
27func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
28func gn(v: i64) -> i64 {
29 let bb: *u8 = sys_mmap(28)
30 var m: i64 = v
31 if m < 0 { sys_write(1, "-\x00" as *u8, 1); m = 0 - m }
32 let t: *u8 = sys_mmap(28)
33 var k: i64 = 0
34 if m == 0 { t[0] = 48 as u8; k = 1 }
35 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
36 var i: i64 = 0
37 while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
38 sys_write(1, bb, k)
39 return 0
40}
41func find_tag(b: *u8, start: i64, end: i64, s: *u8) -> i64 {
42 var i: i64 = start
43 while i + 4 <= end { if mp4_t4(b, i, s) == 1 { return i - 4 } i = i + 1 }
44 return 0 - 1
45}
46
47func main() -> i64 {
48 let szp: *i64 = sys_mmap(16) as *i64
49 let b: *u8 = sys_read_file("knowledge/staging/media/realtest.mp4\x00" as *u8, szp)
50 if (b as i64) == 0 { gp("slice_extract: no realtest.mp4\n\x00" as *u8); return 1 }
51 let len: i64 = szp[0]
52
53 // --- SPS + PPS from avcC ---
54 let avcc: i64 = find_tag(b, 0, len, "avcC\x00" as *u8)
55 if avcc < 0 { gp("slice_extract: no avcC\n\x00" as *u8); return 2 }
56 let pay: i64 = avcc + 8
57 let spsLen: i64 = (b[pay + 6] as i64) * 256 + (b[pay + 7] as i64)
58 let spsdst: *u8 = sys_mmap(spsLen + 16)
59 let spsrbsp: i64 = nx_h264_unescape_rbsp((b as i64 + pay + 9) as *u8, spsLen - 1, spsdst)
60 let sps: *i64 = sys_mmap(128) as *i64
61 nx_h264_parse_sps(spsdst, spsrbsp, sps)
62 let ppspos: i64 = pay + 8 + spsLen
63 let numpps: i64 = b[ppspos] as i64
64 let ppsLen: i64 = (b[ppspos + 1] as i64) * 256 + (b[ppspos + 2] as i64)
65 let ppsdst: *u8 = sys_mmap(ppsLen + 16)
66 let ppsrbsp: i64 = nx_h264_unescape_rbsp((b as i64 + ppspos + 3 + 1) as *u8, ppsLen - 1, ppsdst)
67 let pps: *i64 = sys_mmap(128) as *i64
68 nx_h264_parse_pps(ppsdst, ppsrbsp, pps)
69 gp(" SPS res=\x00" as *u8); gn(sps[2]); gp("x\x00" as *u8); gn(sps[3]); gp(" log2_framenum=\x00" as *u8); gn(sps[9]); gp(" poc_type=\x00" as *u8); gn(sps[10]); gp(" log2_poclsb=\x00" as *u8); gn(sps[11])
70 gp(" | PPS entropy=\x00" as *u8); gn(pps[2]); gp(" qp=\x00" as *u8); gn(pps[3]); gp("\n\x00" as *u8)
71
72 // --- locate sample 0 via video stbl ---
73 let vi: *i64 = sys_mmap(64) as *i64
74 nx_mp4_video_info(b, len, vi)
75 let stbl0: i64 = vi[5]
76 let stbl1: i64 = vi[6]
77 let rco: *i64 = sys_mmap(16) as *i64
78 let rsz: *i64 = sys_mmap(16) as *i64
79 if mp4_find_box(b, stbl0, stbl1, "stco\x00" as *u8, rco) == 0 { gp(" no stco\n\x00" as *u8); return 3 }
80 if mp4_find_box(b, stbl0, stbl1, "stsz\x00" as *u8, rsz) == 0 { gp(" no stsz\n\x00" as *u8); return 4 }
81 let chunk0: i64 = mp4_be32(b, rco[0] + 8) // first chunk file offset
82 let samp_sz_field: i64 = mp4_be32(b, rsz[0] + 4)
83 var samp0: i64 = samp_sz_field
84 if samp_sz_field == 0 { samp0 = mp4_be32(b, rsz[0] + 12) } // per-sample table, sample 0
85 gp(" sample0 file_offset=\x00" as *u8); gn(chunk0); gp(" size=\x00" as *u8); gn(samp0); gp("\n\x00" as *u8)
86
87 // --- NAL-split sample 0 (AVCC, len_size 4) ---
88 let offs: *i64 = sys_mmap(256 * 8) as *i64
89 let types: *i64 = sys_mmap(256 * 8) as *i64
90 let nc: i64 = nx_nal_split_avcc(b, chunk0, chunk0 + samp0, 4, offs, types, 256)
91 gp(" sample0 NALs=\x00" as *u8); gn(nc); gp(" types=[\x00" as *u8)
92 var i: i64 = 0
93 while i < nc { gn(types[i]); gp(" \x00" as *u8); i = i + 1 }
94 gp("]\n\x00" as *u8)
95
96 // --- find the slice NAL (type 1 or 5) and parse its header ---
97 var sidx: i64 = 0 - 1
98 i = 0
99 while i < nc {
100 if types[i] == 5 { if sidx < 0 { sidx = i } }
101 if types[i] == 1 { if sidx < 0 { sidx = i } }
102 i = i + 1
103 }
104 if sidx < 0 { gp(" no slice NAL in sample 0\n\x00" as *u8); return 0 }
105 let noff: i64 = offs[sidx]
106 let ntype: i64 = b[noff] as i64
107 let nlen: i64 = mp4_be32(b, noff - 4)
108 let sl_dst: *u8 = sys_mmap(nlen + 16)
109 let sl_rbsp: i64 = nx_h264_unescape_rbsp((b as i64 + noff + 1) as *u8, nlen - 1, sl_dst)
110 let sh: *i64 = sys_mmap(64) as *i64
111 let sbr: *BitReader = sys_mmap(NX_BR_BYTES) as *BitReader
112 br_init(sbr, sl_dst, sl_rbsp)
113 nx_h264_parse_slice_header_br(sbr, ntype & 31, (ntype >> 5) & 3, sps[9], sps[11], sps[10], sps[6], pps[3], pps[6], sh)
114 gp(" SLICE nal_type=\x00" as *u8); gn(ntype & 31); gp(" len=\x00" as *u8); gn(nlen)
115 gp(" -> first_mb=\x00" as *u8); gn(sh[0]); gp(" slice_type=\x00" as *u8); gn(sh[1]); gp(" (I/P/B=\x00" as *u8); gn(sh[8]); gp(") QP=\x00" as *u8); gn(sh[7]); gp("\n\x00" as *u8)
116
117 // continue into the FIRST macroblock (reader is positioned at slice_data)
118 let mb_type: i64 = nx_h264_mb_parse_itype(sbr)
119 let cls: i64 = nx_h264_mb_type_class(mb_type)
120 gp(" first MB: mb_type=\x00" as *u8); gn(mb_type); gp(" class=\x00" as *u8); gn(cls); gp(" (0=I_NxN 1=I_16x16 2=I_PCM)\x00" as *u8)
121 if cls == 1 {
122 let der: *i64 = sys_mmap(32) as *i64
123 nx_h264_i16x16_derive(mb_type, der)
124 gp(" -> pred=\x00" as *u8); gn(der[0]); gp(" cbpChroma=\x00" as *u8); gn(der[1]); gp(" cbpLuma=\x00" as *u8); gn(der[2])
125 }
126 gp("\n\x00" as *u8)
127
128 if cls == 0 {
129 // I_NxN: pred modes + chroma + CBP, then reach the first residual block
130 let modes: *i64 = sys_mmap(24 * 8) as *i64
131 nx_h264_parse_inxn_predmodes(sbr, modes)
132 let cbp: i64 = nx_h264_decode_cbp_intra(sbr)
133 let cbp_luma: i64 = cbp & 15
134 let cbp_chroma: i64 = (cbp >> 4) & 3
135 gp(" I_NxN chroma_pred=\x00" as *u8); gn(modes[16]); gp(" CBP=\x00" as *u8); gn(cbp)
136 gp(" (luma=\x00" as *u8); gn(cbp_luma); gp(" chroma=\x00" as *u8); gn(cbp_chroma); gp(")\n\x00" as *u8)
137 if (cbp_luma & 1) != 0 {
138 // first luma 4x4 block, nC=0 (no neighbours) -> coeff_token
139 let save_bp: i64 = sbr.byte_pos
140 let save_bit: i64 = sbr.bit_pos
141 let ct: *i64 = sys_mmap(16) as *i64
142 let m: i64 = nx_coeff_token_decode(sbr, ct)
143 if m == 1 {
144 gp(" FIRST RESIDUAL coeff_token(nC=0): TotalCoeff=\x00" as *u8); gn(ct[0]); gp(" TrailingOnes=\x00" as *u8); gn(ct[1]); gp("\n\x00" as *u8)
145 }
146 if m == 0 {
147 sbr.byte_pos = save_bp
148 sbr.bit_pos = save_bit
149 gp(" FIRST RESIDUAL coeff_token: NOT in filled subset; raw bits=\x00" as *u8)
150 var db: i64 = 0
151 while db < 24 { gn(br_read_bit(sbr)); db = db + 1 }
152 gp(" (identify vs Table 9-5 nC<2 -> add entry)\n\x00" as *u8)
153 }
154 }
155 if (cbp_luma & 1) == 0 {
156 gp(" first 8x8 not luma-coded (cbp_luma bit0=0); first residual is chroma/elsewhere\n\x00" as *u8)
157 }
158 }
159
160 if cls == 1 {
161 // I_16x16: mb_pred = intra_chroma_pred_mode (ue) only; CBP derived from mb_type.
162 let chroma_pred: i64 = br_read_ue(sbr)
163 let mb_qp_delta: i64 = br_read_se(sbr)
164 gp(" I_16x16 chroma_pred=\x00" as *u8); gn(chroma_pred); gp(" mb_qp_delta=\x00" as *u8); gn(mb_qp_delta); gp("\n\x00" as *u8)
165 // first residual block = luma DC (always present for I_16x16), nC=0 (no neighbours)
166 let save_bp: i64 = sbr.byte_pos
167 let save_bit: i64 = sbr.bit_pos
168 let ct: *i64 = sys_mmap(16) as *i64
169 let m: i64 = nx_coeff_token_decode(sbr, ct)
170 if m == 1 {
171 gp(" LUMA-DC coeff_token(nC=0): TotalCoeff=\x00" as *u8); gn(ct[0]); gp(" TrailingOnes=\x00" as *u8); gn(ct[1]); gp("\n\x00" as *u8)
172 // decode the levels (algorithmic, no table needed) -> real coefficient values
173 let tc2: i64 = ct[0]
174 let t1b: i64 = ct[1]
175 var slq: i64 = 0
176 if tc2 > 10 { if t1b < 3 { slq = 1 } }
177 let lvl: *i64 = sys_mmap(32 * 8) as *i64
178 var ii: i64 = 0
179 while ii < tc2 {
180 if ii < t1b {
181 let s: i64 = br_read_bit(sbr)
182 if s == 0 { lvl[ii] = 1 }
183 if s == 1 { lvl[ii] = 0 - 1 }
184 }
185 if ii >= t1b {
186 var offq: i64 = 0
187 if ii == t1b { if t1b < 3 { offq = 2 } }
188 let lv: i64 = nx_cavlc_read_level_off(sbr, slq, offq)
189 lvl[ii] = lv
190 if slq == 0 { slq = 1 }
191 var al: i64 = lv
192 if al < 0 { al = 0 - al }
193 if al > (3 << (slq - 1)) { if slq < 6 { slq = slq + 1 } }
194 }
195 ii = ii + 1
196 }
197 gp(" REAL DC levels (reverse-scan):\x00" as *u8)
198 ii = 0
199 while ii < tc2 { gp(" \x00" as *u8); gn(lvl[ii]); ii = ii + 1 }
200 gp("\n\x00" as *u8)
201 // full block via the round-trip-proven assembler (rewind to block start)
202 sbr.byte_pos = save_bp
203 sbr.bit_pos = save_bit
204 let dcblk: *i64 = sys_mmap(16 * 8) as *i64
205 let dtc: i64 = nx_h264_residual_decode(sbr, 16, 0, dcblk)
206 gp(" REAL DC block (scan order, n=\x00" as *u8); gn(dtc); gp("):\x00" as *u8)
207 var jj: i64 = 0
208 while jj < 16 { gp(" \x00" as *u8); gn(dcblk[jj]); jj = jj + 1 }
209 gp(" <- first real block fully decoded from realtest.mp4\n\x00" as *u8)
210
211 // ===== I_16x16 luma RECONSTRUCT (cbpLuma=0 -> DC-only) + compare vs ffmpeg =====
212 let mbqp: i64 = sh[7] // MB QP (slice QP + mb_qp_delta 0) = 31
213 let dcras: *i64 = sys_mmap(16 * 8) as *i64
214 nx_h264_inv_zigzag4x4(dcblk, dcras) // DC scan -> raster (one DC per 4x4 sub-block)
215 nx_h264_hadamard4x4(dcras) // inverse Hadamard
216 let mm2: i64 = mbqp % 6
217 let sh6: i64 = mbqp / 6
218 let na: *i64 = sys_mmap(8 * 8) as *i64
219 na[0]=10; na[1]=11; na[2]=13; na[3]=14; na[4]=16; na[5]=18
220 let ls: i64 = 16 * na[mm2]
221 let dcy: *i64 = sys_mmap(16 * 8) as *i64
222 var qi: i64 = 0
223 while qi < 16 {
224 var v: i64 = 0
225 if mbqp >= 36 { v = (dcras[qi] * ls) << (sh6 - 6) }
226 if mbqp < 36 { v = (dcras[qi] * ls + (1 << (5 - sh6))) >> (6 - sh6) }
227 dcy[qi] = v
228 qi = qi + 1
229 }
230 let recon: *u8 = sys_mmap(256)
231 let blk: *i64 = sys_mmap(16 * 8) as *i64
232 var br2: i64 = 0
233 while br2 < 4 {
234 var bc: i64 = 0
235 while bc < 4 {
236 var z: i64 = 0
237 while z < 16 { blk[z] = 0; z = z + 1 }
238 blk[0] = dcy[br2 * 4 + bc]
239 nx_idct4x4(blk) // DC-only -> flat residual
240 let res: i64 = blk[0]
241 var py: i64 = 0
242 while py < 4 {
243 var px: i64 = 0
244 while px < 4 {
245 var val: i64 = 128 + res // Intra_16x16 DC pred (first MB, no neighbours = 128)
246 if val < 0 { val = 0 }
247 if val > 255 { val = 255 }
248 recon[(br2 * 4 + py) * 16 + (bc * 4 + px)] = val as u8
249 px = px + 1
250 }
251 py = py + 1
252 }
253 bc = bc + 1
254 }
255 br2 = br2 + 1
256 }
257 // ===== in-loop DEBLOCKING: ALL internal vertical edges x=4,8,12 in order (bS=3) =====
258 // (cascade matters: x=8's output feeds x=12, which then corrects the q2 column)
259 let vedges: *i64 = sys_mmap(8 * 8) as *i64
260 vedges[0] = 4; vedges[1] = 8; vedges[2] = 12
261 let edge: *i64 = sys_mmap(8 * 8) as *i64
262 var ei: i64 = 0
263 while ei < 3 {
264 let ex: i64 = vedges[ei]
265 var ry: i64 = 0
266 while ry < 16 {
267 edge[0] = recon[ry * 16 + ex - 4] as i64; edge[1] = recon[ry * 16 + ex - 3] as i64
268 edge[2] = recon[ry * 16 + ex - 2] as i64; edge[3] = recon[ry * 16 + ex - 1] as i64
269 edge[4] = recon[ry * 16 + ex] as i64; edge[5] = recon[ry * 16 + ex + 1] as i64
270 edge[6] = recon[ry * 16 + ex + 2] as i64; edge[7] = recon[ry * 16 + ex + 3] as i64
271 nx_deblock_luma_edge(edge, 3, mbqp)
272 recon[ry * 16 + ex - 2] = edge[2] as u8; recon[ry * 16 + ex - 1] = edge[3] as u8
273 recon[ry * 16 + ex] = edge[4] as u8; recon[ry * 16 + ex + 1] = edge[5] as u8
274 ry = ry + 1
275 }
276 ei = ei + 1
277 }
278 gp(" RECON 16x16 luma (post-deblock) topleft=\x00" as *u8); gn(recon[0] as i64); gp(" mid=\x00" as *u8); gn(recon[8] as i64); gp(" topright=\x00" as *u8); gn(recon[15] as i64); gp("\n\x00" as *u8)
279 // compare against ffmpeg ref_frame0_gray.pgm top-left 16x16 (P5 header 16 bytes, width 576)
280 // compare against ffmpeg RAW Y plane (limited-range, no header, width 576) -- the true ref
281 let rzp: *i64 = sys_mmap(16) as *i64
282 let refp: *u8 = sys_read_file("knowledge/staging/media/ref_frame0.yuv\x00" as *u8, rzp)
283 if (refp as i64) != 0 {
284 var maxd: i64 = 0
285 var sumd: i64 = 0
286 var nexact: i64 = 0
287 var yy: i64 = 0
288 while yy < 16 {
289 var xx: i64 = 0
290 while xx < 16 {
291 let mine: i64 = recon[yy * 16 + xx] as i64
292 let theirs: i64 = refp[yy * 576 + xx] as i64
293 var dd: i64 = mine - theirs
294 if dd < 0 { dd = 0 - dd }
295 if dd > maxd { maxd = dd }
296 if dd == 0 { nexact = nexact + 1 }
297 sumd = sumd + dd
298 xx = xx + 1
299 }
300 yy = yy + 1
301 }
302 gp(" vs ffmpeg RAW Y (post-deblock): exact=\x00" as *u8); gn(nexact); gp("/256 MAXDIFF=\x00" as *u8); gn(maxd); gp(" meandiff_x100=\x00" as *u8); gn(sumd * 100 / 256)
303 if maxd == 0 { gp(" *** BIT-EXACT vs ffmpeg -- full per-MB sovereign decode VERIFIED ***\n\x00" as *u8) }
304 if maxd != 0 { gp(" (close; residual nuance)\n\x00" as *u8) }
305 }
306
307 // ===== finish MB0: chroma (cbpChroma -> DC only when ==1, U then V) + reach MB1 =====
308 let cbpc: i64 = ((mb_type - 1) / 4) % 3 // I_16x16 chroma CBP derived from mb_type
309 if cbpc >= 1 {
310 let ucdc: *i64 = sys_mmap(8 * 8) as *i64
311 let utc: i64 = nx_h264_residual_decode_cdc(sbr, ucdc)
312 let vcdc: *i64 = sys_mmap(8 * 8) as *i64
313 let vtc: i64 = nx_h264_residual_decode_cdc(sbr, vcdc)
314 gp(" chroma DC: U tc=\x00" as *u8); gn(utc); gp(" [\x00" as *u8); gn(ucdc[0]); gp(" \x00" as *u8); gn(ucdc[1]); gp(" \x00" as *u8); gn(ucdc[2]); gp(" \x00" as *u8); gn(ucdc[3]); gp("] V tc=\x00" as *u8); gn(vtc); gp(" [\x00" as *u8); gn(vcdc[0]); gp(" \x00" as *u8); gn(vcdc[1]); gp(" \x00" as *u8); gn(vcdc[2]); gp(" \x00" as *u8); gn(vcdc[3]); gp("]\n\x00" as *u8)
315 if cbpc == 1 {
316 // cbpChroma==1: DC only, no chroma AC -> MB0 fully parsed; read MB1 mb_type
317 let mb1: i64 = nx_h264_mb_parse_itype(sbr)
318 let cls1: i64 = nx_h264_mb_type_class(mb1)
319 gp(" -> MB0 FULLY PARSED (luma+chroma). MB1 mb_type=\x00" as *u8); gn(mb1); gp(" class=\x00" as *u8); gn(cls1)
320 if cls1 == 0 { gp(" I_NxN\x00" as *u8) }
321 if cls1 == 1 { gp(" I_16x16\x00" as *u8) }
322 if cls1 == 2 { gp(" I_PCM\x00" as *u8) }
323 if mb1 < 0 { gp(" (INVALID -> chroma desync)\x00" as *u8) }
324 if mb1 > 25 { gp(" (INVALID -> chroma desync)\x00" as *u8) }
325 gp("\n\x00" as *u8)
326 }
327 }
328 }
329 if m == 0 {
330 sbr.byte_pos = save_bp
331 sbr.bit_pos = save_bit
332 gp(" LUMA-DC coeff_token NOT in filled subset; raw bits=\x00" as *u8)
333 var db: i64 = 0
334 while db < 24 { gn(br_read_bit(sbr)); db = db + 1 }
335 gp(" (identify vs Table 9-5 nC<2 -> add entry)\n\x00" as *u8)
336 }
337 }
338 return 0
339}