]>
Commit | Line | Data |
---|---|---|
1 | use nihav_core::frame::*; | |
2 | use nihav_codec_support::codecs::{MV, ZERO_MV}; | |
3 | use super::super::vpcommon::*; | |
4 | use super::VP56DCPred; | |
5 | use super::dsp::*; | |
6 | use super::rdo::*; | |
7 | ||
8 | /*#[cfg(debug_assertions)] | |
9 | use std::io::Write; | |
10 | #[cfg(debug_assertions)] | |
11 | use std::fs::File; | |
12 | #[cfg(debug_assertions)] | |
13 | pub fn dump_pgm(vbuf: &NAVideoBuffer<u8>, name: &str) { | |
14 | let dst = vbuf.get_data(); | |
15 | let (w, h) = vbuf.get_dimensions(0); | |
16 | let mut file = File::create(name).unwrap(); | |
17 | file.write_all(format!("P5\n{} {}\n255\n", w, h * 3 / 2).as_bytes()).unwrap(); | |
18 | for row in dst[vbuf.get_offset(0)..].chunks(vbuf.get_stride(0)).take(h).rev() { | |
19 | file.write_all(row).unwrap(); | |
20 | } | |
21 | for (row1, row2) in dst[vbuf.get_offset(1)..].chunks(vbuf.get_stride(1)).take(h / 2).zip(dst[vbuf.get_offset(2)..].chunks(vbuf.get_stride(2))).rev() { | |
22 | file.write_all(row1).unwrap(); | |
23 | file.write_all(row2).unwrap(); | |
24 | } | |
25 | }*/ | |
26 | ||
27 | pub type Coeffs = [[i16; 64]; 6]; | |
28 | ||
29 | #[derive(Clone)] | |
30 | pub struct ResidueMB { | |
31 | pub coeffs: Coeffs, | |
32 | } | |
33 | ||
34 | impl ResidueMB { | |
35 | fn new() -> Self { | |
36 | Self { | |
37 | coeffs: [[0; 64]; 6], | |
38 | } | |
39 | } | |
40 | fn fdct(&mut self) { | |
41 | for blk in self.coeffs.iter_mut() { | |
42 | vp_fdct(blk); | |
43 | } | |
44 | } | |
45 | fn idct(&mut self) { | |
46 | for blk in self.coeffs.iter_mut() { | |
47 | vp_idct(blk); | |
48 | } | |
49 | } | |
50 | fn quant(&mut self, q: usize) { | |
51 | for blk in self.coeffs.iter_mut() { | |
52 | if blk[0] != 0 { | |
53 | blk[0] /= VP56_DC_QUANTS[q] * 4; | |
54 | } | |
55 | for coef in blk[1..].iter_mut() { | |
56 | if *coef != 0 { | |
57 | *coef /= VP56_AC_QUANTS[q] * 4; | |
58 | } | |
59 | } | |
60 | } | |
61 | } | |
62 | fn dequant(&mut self, q: usize) { | |
63 | for blk in self.coeffs.iter_mut() { | |
64 | if blk[0] != 0 { | |
65 | blk[0] *= VP56_DC_QUANTS[q] * 4; | |
66 | } | |
67 | for coef in blk[1..].iter_mut() { | |
68 | if *coef != 0 { | |
69 | *coef *= VP56_AC_QUANTS[q] * 4; | |
70 | } | |
71 | } | |
72 | } | |
73 | } | |
74 | fn dequant_from(&mut self, src: &Self, q: usize) { | |
75 | for (dblk, sblk) in self.coeffs.iter_mut().zip(src.coeffs.iter()) { | |
76 | dblk[0] = if sblk[0] != 0 { sblk[0] * VP56_DC_QUANTS[q] * 4 } else { 0 }; | |
77 | for (dcoef, &scoef) in dblk[1..].iter_mut().zip(sblk[1..].iter()) { | |
78 | *dcoef = if scoef != 0 { scoef * VP56_AC_QUANTS[q] * 4 } else { 0 }; | |
79 | } | |
80 | } | |
81 | } | |
82 | fn fill(&self, dst: &mut [[u8; 64]; 6]) { | |
83 | for (dblk, sblk) in dst.iter_mut().zip(self.coeffs.iter()) { | |
84 | for (dcoef, &scoef) in dblk.iter_mut().zip(sblk.iter()) { | |
85 | *dcoef = scoef as u8; | |
86 | } | |
87 | } | |
88 | } | |
89 | } | |
90 | ||
91 | #[derive(Clone)] | |
92 | pub struct InterMB { | |
93 | pub residue: ResidueMB, | |
94 | pub reference: Coeffs, | |
95 | pub mv: [MV; 4], | |
96 | } | |
97 | ||
98 | impl InterMB { | |
99 | fn new() -> Self { | |
100 | Self { | |
101 | residue: ResidueMB::new(), | |
102 | reference: [[0; 64]; 6], | |
103 | mv: [ZERO_MV; 4], | |
104 | } | |
105 | } | |
106 | } | |
107 | ||
108 | const VP56_DC_QUANTS: [i16; 64] = [ | |
109 | 47, 47, 47, 47, 45, 43, 43, 43, | |
110 | 43, 43, 42, 41, 41, 40, 40, 40, | |
111 | 40, 35, 35, 35, 35, 33, 33, 33, | |
112 | 33, 32, 32, 32, 27, 27, 26, 26, | |
113 | 25, 25, 24, 24, 23, 23, 19, 19, | |
114 | 19, 19, 18, 18, 17, 16, 16, 16, | |
115 | 16, 16, 15, 11, 11, 11, 10, 10, | |
116 | 9, 8, 7, 5, 3, 3, 2, 2 | |
117 | ]; | |
118 | const VP56_AC_QUANTS: [i16; 64] = [ | |
119 | 94, 92, 90, 88, 86, 82, 78, 74, | |
120 | 70, 66, 62, 58, 54, 53, 52, 51, | |
121 | 50, 49, 48, 47, 46, 45, 44, 43, | |
122 | 42, 40, 39, 37, 36, 35, 34, 33, | |
123 | 32, 31, 30, 29, 28, 27, 26, 25, | |
124 | 24, 23, 22, 21, 20, 19, 18, 17, | |
125 | 16, 15, 14, 13, 12, 11, 10, 9, | |
126 | 8, 7, 6, 5, 4, 3, 2, 1 | |
127 | ]; | |
128 | ||
129 | const VP56_FILTER_LIMITS: [u8; 64] = [ | |
130 | 14, 14, 13, 13, 12, 12, 10, 10, | |
131 | 10, 10, 8, 8, 8, 8, 8, 8, | |
132 | 8, 8, 8, 8, 8, 8, 8, 8, | |
133 | 8, 8, 8, 8, 8, 8, 8, 8, | |
134 | 8, 8, 8, 8, 7, 7, 7, 7, | |
135 | 7, 7, 6, 6, 6, 6, 6, 6, | |
136 | 5, 5, 5, 5, 4, 4, 4, 4, | |
137 | 4, 4, 4, 3, 3, 3, 3, 2 | |
138 | ]; | |
139 | ||
140 | #[derive(Default)] | |
141 | pub struct FrameEncoder { | |
142 | pub quant: usize, | |
143 | pub src_mbs: Vec<ResidueMB>, | |
144 | pub intra_mbs: Vec<ResidueMB>, | |
145 | pub inter_mbs: Vec<InterMB>, | |
146 | pub fourmv_mbs: Vec<InterMB>, | |
147 | pub golden_mbs: Vec<InterMB>, | |
148 | ||
149 | pub mb_types: Vec<VPMBType>, | |
150 | pub num_mv: Vec<u8>, | |
151 | pub coded_mv: Vec<[MV; 4]>, | |
152 | pub fmv_sub: Vec<[VPMBType; 4]>, | |
153 | ||
154 | pub mb_w: usize, | |
155 | pub mb_h: usize, | |
156 | ||
157 | pub me_mode: MVSearchMode, | |
158 | pub me_range: i16, | |
159 | } | |
160 | ||
161 | macro_rules! read_block { | |
162 | ($dst: expr, $src: expr, $stride: expr) => { | |
163 | for (drow, srow) in $dst.chunks_mut(8).zip($src.chunks($stride).take(8)) { | |
164 | for (dst, &src) in drow.iter_mut().zip(srow.iter()) { | |
165 | *dst = i16::from(src); | |
166 | } | |
167 | } | |
168 | } | |
169 | } | |
170 | ||
171 | macro_rules! write_block { | |
172 | ($dst: expr, $src: expr, $stride: expr) => { | |
173 | for (drow, srow) in $dst.chunks_mut($stride).take(8).zip($src.chunks(8)) { | |
174 | drow[..8].copy_from_slice(srow); | |
175 | } | |
176 | } | |
177 | } | |
178 | ||
179 | impl FrameEncoder { | |
180 | pub fn new() -> Self { Self::default() } | |
181 | pub fn resize(&mut self, mb_w: usize, mb_h: usize) { | |
182 | self.mb_w = mb_w; | |
183 | self.mb_h = mb_h; | |
184 | ||
185 | let num_mbs = self.mb_w * self.mb_h; | |
186 | self.src_mbs.clear(); | |
187 | self.src_mbs.reserve(num_mbs); | |
188 | self.intra_mbs.clear(); | |
189 | self.intra_mbs.reserve(num_mbs); | |
190 | self.inter_mbs.clear(); | |
191 | self.inter_mbs.reserve(num_mbs); | |
192 | self.fourmv_mbs.clear(); | |
193 | self.fourmv_mbs.reserve(num_mbs); | |
194 | self.golden_mbs.clear(); | |
195 | self.golden_mbs.reserve(num_mbs); | |
196 | ||
197 | self.mb_types.clear(); | |
198 | self.mb_types.reserve(num_mbs); | |
199 | self.num_mv.clear(); | |
200 | self.num_mv.reserve(num_mbs); | |
201 | self.coded_mv.clear(); | |
202 | self.coded_mv.reserve(num_mbs); | |
203 | self.fmv_sub.clear(); | |
204 | self.fmv_sub.reserve(num_mbs); | |
205 | } | |
206 | pub fn set_quant(&mut self, quant: usize) { self.quant = quant; } | |
207 | pub fn read_mbs(&mut self, vbuf: &NAVideoBuffer<u8>) { | |
208 | let src = vbuf.get_data(); | |
209 | let y = &src[vbuf.get_offset(0)..]; | |
210 | let ystride = vbuf.get_stride(0); | |
211 | let u = &src[vbuf.get_offset(1)..]; | |
212 | let ustride = vbuf.get_stride(1); | |
213 | let v = &src[vbuf.get_offset(2)..]; | |
214 | let vstride = vbuf.get_stride(2); | |
215 | let (w, _) = vbuf.get_dimensions(0); | |
216 | ||
217 | self.src_mbs.clear(); | |
218 | for (ys, (us, vs)) in y.chunks(ystride * 16).zip(u.chunks(ustride * 8).zip(v.chunks(vstride * 8))) { | |
219 | for x in (0..w).step_by(16) { | |
220 | let mut mb = ResidueMB::new(); | |
221 | for (i, blk) in mb.coeffs[..4].iter_mut().enumerate() { | |
222 | read_block!(blk, ys[x + (i & 1) * 8 + (i >> 1) * 8 * ystride..], ystride); | |
223 | } | |
224 | read_block!(mb.coeffs[4], us[x/2..], ustride); | |
225 | read_block!(mb.coeffs[5], vs[x/2..], vstride); | |
226 | self.src_mbs.push(mb); | |
227 | } | |
228 | } | |
229 | } | |
230 | pub fn reconstruct_frame(&mut self, dc_pred: &mut VP56DCPred, mut vbuf: NAVideoBufferRef<u8>) { | |
231 | let mut blocks = [[0u8; 64]; 6]; | |
232 | ||
233 | let mut yoff = vbuf.get_offset(0); | |
234 | let mut uoff = vbuf.get_offset(1); | |
235 | let mut voff = vbuf.get_offset(2); | |
236 | let ystride = vbuf.get_stride(0); | |
237 | let ustride = vbuf.get_stride(1); | |
238 | let vstride = vbuf.get_stride(2); | |
239 | let dst = vbuf.get_data_mut().unwrap(); | |
240 | ||
241 | dc_pred.reset(); | |
242 | ||
243 | let quant = self.quant; | |
244 | let mut mb_pos = 0; | |
245 | for _mb_y in 0..self.mb_h { | |
246 | for mb_x in 0..self.mb_w { | |
247 | let mb_type = self.mb_types[mb_pos]; | |
248 | let mb = self.get_mb_mut(mb_pos); | |
249 | for (i, blk) in mb.coeffs.iter_mut().enumerate() { | |
250 | dc_pred.predict_dc(mb_type, i, blk, false); | |
251 | } | |
252 | mb.dequant(quant); | |
253 | mb.idct(); | |
254 | let mb = self.get_mb(mb_pos); | |
255 | if mb_type.is_intra() { | |
256 | for (dblk, sblk) in blocks.iter_mut().zip(mb.coeffs.iter()) { | |
257 | for (dcoef, &scoef) in dblk.iter_mut().zip(sblk.iter()) { | |
258 | *dcoef = (scoef + 128).max(0).min(255) as u8; | |
259 | } | |
260 | } | |
261 | } else { | |
262 | let res_mb = match mb_type.get_ref_id() { | |
263 | 0 => unreachable!(), | |
264 | 1 => if mb_type != VPMBType::InterFourMV { | |
265 | &self.inter_mbs[mb_pos].reference | |
266 | } else { | |
267 | &self.fourmv_mbs[mb_pos].reference | |
268 | }, | |
269 | _ => &self.golden_mbs[mb_pos].reference, | |
270 | }; | |
271 | ||
272 | for (dblk, (sblk1, sblk2)) in blocks.iter_mut().zip(mb.coeffs.iter().zip(res_mb.iter())) { | |
273 | for (dcoef, (&scoef1, &scoef2)) in dblk.iter_mut().zip(sblk1.iter().zip(sblk2.iter())) { | |
274 | *dcoef = (scoef1 + scoef2).max(0).min(255) as u8; | |
275 | } | |
276 | } | |
277 | } | |
278 | ||
279 | for i in 0..4 { | |
280 | write_block!(&mut dst[yoff + mb_x * 16 + (i & 1) * 8 + (i >> 1) * 8 * ystride..], | |
281 | blocks[i], ystride); | |
282 | } | |
283 | write_block!(&mut dst[uoff + mb_x * 8..], blocks[4], ustride); | |
284 | write_block!(&mut dst[voff + mb_x * 8..], blocks[5], vstride); | |
285 | ||
286 | dc_pred.next_mb(); | |
287 | mb_pos += 1; | |
288 | } | |
289 | yoff += ystride * 16; | |
290 | uoff += ustride * 8; | |
291 | voff += vstride * 8; | |
292 | dc_pred.update_row(); | |
293 | } | |
294 | /*#[cfg(debug_assertions)] | |
295 | dump_pgm(&vbuf, "/home/kst/devel/NihAV-rust/assets/test_out/debug.pgm");*/ | |
296 | } | |
297 | pub fn get_mb(&self, mb_pos: usize) -> &ResidueMB { | |
298 | let mb_type = self.mb_types[mb_pos]; | |
299 | match mb_type.get_ref_id() { | |
300 | 0 => &self.intra_mbs[mb_pos], | |
301 | 1 => if mb_type != VPMBType::InterFourMV { | |
302 | &self.inter_mbs[mb_pos].residue | |
303 | } else { | |
304 | &self.fourmv_mbs[mb_pos].residue | |
305 | }, | |
306 | _ => &self.golden_mbs[mb_pos].residue, | |
307 | } | |
308 | } | |
309 | fn get_mb_mut(&mut self, mb_pos: usize) -> &mut ResidueMB { | |
310 | let mb_type = self.mb_types[mb_pos]; | |
311 | match mb_type.get_ref_id() { | |
312 | 0 => &mut self.intra_mbs[mb_pos], | |
313 | 1 => if mb_type != VPMBType::InterFourMV { | |
314 | &mut self.inter_mbs[mb_pos].residue | |
315 | } else { | |
316 | &mut self.fourmv_mbs[mb_pos].residue | |
317 | }, | |
318 | _ => &mut self.golden_mbs[mb_pos].residue, | |
319 | } | |
320 | } | |
321 | pub fn prepare_intra_blocks(&mut self) { | |
322 | self.intra_mbs.clear(); | |
323 | self.mb_types.clear(); | |
324 | for smb in self.src_mbs.iter() { | |
325 | let mut dmb = smb.clone(); | |
326 | dmb.fdct(); | |
327 | for blk in dmb.coeffs.iter_mut() { | |
328 | blk[0] -= 4096; | |
329 | } | |
330 | dmb.quant(self.quant); | |
331 | self.mb_types.push(VPMBType::Intra); | |
332 | self.intra_mbs.push(dmb); | |
333 | } | |
334 | } | |
335 | pub fn prepare_inter_blocks(&mut self, golden: bool) { | |
336 | let inter_mbs = if !golden { &mut self.inter_mbs } else { &mut self.golden_mbs }; | |
337 | for (mb_idx, mb) in inter_mbs.iter_mut().enumerate() { | |
338 | mb.residue.fdct(); | |
339 | mb.residue.quant(self.quant); | |
340 | self.mb_types[mb_idx] = VPMBType::InterMV; | |
341 | } | |
342 | } | |
343 | pub fn estimate_mvs(&mut self, ref_frame: NAVideoBufferRef<u8>, mc_buf: NAVideoBufferRef<u8>, golden: bool) { | |
344 | let loop_thr = i16::from(VP56_FILTER_LIMITS[self.quant as usize]); | |
345 | ||
346 | let inter_mbs = if !golden { &mut self.inter_mbs } else { &mut self.golden_mbs }; | |
347 | ||
348 | if inter_mbs.is_empty() { | |
349 | for _ in 0..self.mb_w * self.mb_h { | |
350 | inter_mbs.push(InterMB::new()); | |
351 | } | |
352 | } | |
353 | ||
354 | let mut cur_blk = [[0u8; 64]; 6]; | |
355 | ||
356 | let mut mv_est = MVEstimator::new(ref_frame, mc_buf, loop_thr, self.me_range); | |
357 | ||
358 | let mut mv_search = self.me_mode.create_search(); | |
359 | ||
360 | let mut mb_pos = 0; | |
361 | for (mb_y, row) in inter_mbs.chunks_mut(self.mb_w).enumerate() { | |
362 | for (mb_x, mb) in row.iter_mut().enumerate() { | |
363 | self.src_mbs[mb_pos].fill(&mut cur_blk); | |
364 | ||
365 | let (best_mv, _best_dist) = mv_search.search_mb(&mut mv_est, &cur_blk, mb_x, mb_y); | |
366 | mb.mv[3] = best_mv; | |
367 | ||
368 | for i in 0..4 { | |
369 | mv_est.mc_block(i, 0, mb_x * 16 + (i & 1) * 8, mb_y * 16 + (i >> 1) * 8, best_mv); | |
370 | sub_blk(&mut mb.residue.coeffs[i], &cur_blk[i], &mv_est.ref_blk[i]); | |
371 | } | |
372 | for plane in 1..3 { | |
373 | mv_est.mc_block(plane + 3, plane, mb_x * 8, mb_y * 8, best_mv); | |
374 | sub_blk(&mut mb.residue.coeffs[plane + 3], &cur_blk[plane + 3], &mv_est.ref_blk[plane + 3]); | |
375 | } | |
376 | ||
377 | for (dblk, sblk) in mb.reference.iter_mut().zip(mv_est.ref_blk.iter()) { | |
378 | for (dst, &src) in dblk.iter_mut().zip(sblk.iter()) { | |
379 | *dst = i16::from(src); | |
380 | } | |
381 | } | |
382 | mb_pos += 1; | |
383 | } | |
384 | } | |
385 | } | |
386 | fn estimate_fourmv(&mut self, ref_frame: NAVideoBufferRef<u8>, mc_buf: NAVideoBufferRef<u8>, mb_x: usize, mb_y: usize) -> bool { | |
387 | let loop_thr = i16::from(VP56_FILTER_LIMITS[self.quant as usize]); | |
388 | ||
389 | if self.fourmv_mbs.is_empty() { | |
390 | for _ in 0..self.mb_w * self.mb_h { | |
391 | self.fourmv_mbs.push(InterMB::new()); | |
392 | } | |
393 | } | |
394 | if self.fmv_sub.is_empty() { | |
395 | self.fmv_sub.resize(self.mb_w * self.mb_h, [VPMBType::Intra; 4]); | |
396 | } | |
397 | ||
398 | let mb_pos = mb_x + mb_y * self.mb_w; | |
399 | let mb = &mut self.fourmv_mbs[mb_pos]; | |
400 | ||
401 | let mut cur_blk = [[0u8; 64]; 6]; | |
402 | self.src_mbs[mb_pos].fill(&mut cur_blk); | |
403 | ||
404 | let mut mv_est = MVEstimator::new(ref_frame, mc_buf, loop_thr, self.me_range); | |
405 | ||
406 | let mut mv_search = self.me_mode.create_search(); | |
407 | ||
408 | for i in 0..4 { | |
409 | let xpos = mb_x * 16 + (i & 1) * 8; | |
410 | let ypos = mb_y * 16 + (i >> 1) * 8; | |
411 | let (best_mv, _best_dist) = mv_search.search_blk(&mut mv_est, &cur_blk[i], xpos, ypos); | |
412 | mb.mv[i] = best_mv; | |
413 | } | |
414 | let mvsum = mb.mv[0] + mb.mv[1] + mb.mv[2] + mb.mv[3]; | |
415 | let chroma_mv = MV{ x: mvsum.x / 4, y: mvsum.y / 4}; | |
416 | ||
417 | for (i, blk) in mb.residue.coeffs[..4].iter_mut().enumerate() { | |
418 | let xpos = mb_x * 16 + (i & 1) * 8; | |
419 | let ypos = mb_y * 16 + (i >> 1) * 8; | |
420 | mv_est.mc_block(i, 0, xpos, ypos, mb.mv[i]); | |
421 | sub_blk(blk, &cur_blk[i], &mv_est.ref_blk[i]); | |
422 | } | |
423 | for plane in 1..3 { | |
424 | mv_est.mc_block(plane + 3, plane, mb_x * 8, mb_y * 8, chroma_mv); | |
425 | sub_blk(&mut mb.residue.coeffs[plane + 3], &cur_blk[plane + 3], &mv_est.ref_blk[plane + 3]); | |
426 | } | |
427 | ||
428 | for (dblk, sblk) in mb.reference.iter_mut().zip(mv_est.ref_blk.iter()) { | |
429 | for (dst, &src) in dblk.iter_mut().zip(sblk.iter()) { | |
430 | *dst = i16::from(src); | |
431 | } | |
432 | } | |
433 | ||
434 | (mb.mv[0] != mb.mv[1]) || (mb.mv[0] != mb.mv[2]) || (mb.mv[0] != mb.mv[3]) | |
435 | } | |
436 | pub fn select_inter_blocks(&mut self, ref_frame: NAVideoBufferRef<u8>, mc_buf: NAVideoBufferRef<u8>, has_golden_frame: bool, lambda: f32) { | |
437 | let mut tmp_mb = ResidueMB::new(); | |
438 | for mb_idx in 0..self.mb_w * self.mb_h { | |
439 | tmp_mb.dequant_from(&self.intra_mbs[mb_idx], self.quant); | |
440 | tmp_mb.idct(); | |
441 | for blk in tmp_mb.coeffs.iter_mut() { | |
442 | for coef in blk.iter_mut() { | |
443 | *coef = (*coef + 128).max(0).min(255); | |
444 | } | |
445 | } | |
446 | let intra_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
447 | let intra_nits = estimate_intra_mb_nits(&self.intra_mbs[mb_idx].coeffs, self.quant); | |
448 | let intra_cost = (intra_dist as f32) + lambda * (intra_nits as f32); | |
449 | ||
450 | tmp_mb.dequant_from(&self.inter_mbs[mb_idx].residue, self.quant); | |
451 | tmp_mb.idct(); | |
452 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.inter_mbs[mb_idx].reference.iter()) { | |
453 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
454 | *coef = (*coef + add).max(0).min(255); | |
455 | } | |
456 | } | |
457 | let inter_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
458 | let mut inter_nits = estimate_inter_mb_nits(&self.inter_mbs[mb_idx], self.quant, false); | |
459 | if self.inter_mbs[mb_idx].mv[3] != ZERO_MV { | |
460 | inter_nits += estimate_mv_nits(self.inter_mbs[mb_idx].mv[3]); | |
461 | } | |
462 | let mut inter_cost = (inter_dist as f32) + lambda * (inter_nits as f32); | |
463 | ||
464 | if inter_cost < intra_cost { | |
465 | self.mb_types[mb_idx] = VPMBType::InterMV; | |
466 | ||
467 | if inter_dist > 512 { | |
468 | self.estimate_fourmv(ref_frame.clone(), mc_buf.clone(), mb_idx % self.mb_w, mb_idx / self.mb_w); | |
469 | self.fourmv_mbs[mb_idx].residue.fdct(); | |
470 | self.fourmv_mbs[mb_idx].residue.quant(self.quant); | |
471 | ||
472 | tmp_mb.dequant_from(&self.fourmv_mbs[mb_idx].residue, self.quant); | |
473 | tmp_mb.idct(); | |
474 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.fourmv_mbs[mb_idx].reference.iter()) { | |
475 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
476 | *coef = (*coef + add).max(0).min(255); | |
477 | } | |
478 | } | |
479 | let fourmv_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
480 | let fourmv_nits = estimate_inter_mb_nits(&self.fourmv_mbs[mb_idx], self.quant, true); | |
481 | let fourmv_cost = (fourmv_dist as f32) + lambda * (fourmv_nits as f32); | |
482 | if fourmv_cost < inter_cost { | |
483 | self.mb_types[mb_idx] = VPMBType::InterFourMV; | |
484 | inter_cost = fourmv_cost; | |
485 | } | |
486 | } | |
487 | } | |
488 | ||
489 | if has_golden_frame { | |
490 | tmp_mb.dequant_from(&self.golden_mbs[mb_idx].residue, self.quant); | |
491 | tmp_mb.idct(); | |
492 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.golden_mbs[mb_idx].reference.iter()) { | |
493 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
494 | *coef = (*coef + add).max(0).min(255); | |
495 | } | |
496 | } | |
497 | let golden_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
498 | let golden_nits = estimate_inter_mb_nits(&self.golden_mbs[mb_idx], self.quant, false); | |
499 | let golden_cost = (golden_dist as f32) + lambda * (golden_nits as f32); | |
500 | ||
501 | if (self.mb_types[mb_idx].is_intra() && golden_cost < intra_cost) || | |
502 | (!self.mb_types[mb_idx].is_intra() && golden_cost < inter_cost) { | |
503 | self.mb_types[mb_idx] = VPMBType::GoldenMV; | |
504 | } | |
505 | } | |
506 | } | |
507 | } | |
508 | fn motion_est_mb(src_mb: &ResidueMB, cur_blk: &mut [[u8; 64]; 6], mb: &mut InterMB, mv_search: &mut Box<dyn MVSearch+Send>, mv_est: &mut MVEstimator, mb_x: usize, mb_y: usize) { | |
509 | src_mb.fill(cur_blk); | |
510 | let (best_mv, _best_dist) = mv_search.search_mb(mv_est, cur_blk, mb_x, mb_y); | |
511 | mb.mv[3] = best_mv; | |
512 | ||
513 | for i in 0..4 { | |
514 | mv_est.mc_block(i, 0, mb_x * 16 + (i & 1) * 8, mb_y * 16 + (i >> 1) * 8, best_mv); | |
515 | sub_blk(&mut mb.residue.coeffs[i], &cur_blk[i], &mv_est.ref_blk[i]); | |
516 | } | |
517 | for plane in 1..3 { | |
518 | mv_est.mc_block(plane + 3, plane, mb_x * 8, mb_y * 8, best_mv); | |
519 | sub_blk(&mut mb.residue.coeffs[plane + 3], &cur_blk[plane + 3], &mv_est.ref_blk[plane + 3]); | |
520 | } | |
521 | ||
522 | for (dblk, sblk) in mb.reference.iter_mut().zip(mv_est.ref_blk.iter()) { | |
523 | for (dst, &src) in dblk.iter_mut().zip(sblk.iter()) { | |
524 | *dst = i16::from(src); | |
525 | } | |
526 | } | |
527 | } | |
528 | pub fn select_inter_blocks_fast(&mut self, ref_frame: NAVideoBufferRef<u8>, gold_frame: Option<NAVideoBufferRef<u8>>, mc_buf: NAVideoBufferRef<u8>, lambda: f32) { | |
529 | let loop_thr = i16::from(VP56_FILTER_LIMITS[self.quant as usize]); | |
530 | ||
531 | if self.inter_mbs.is_empty() { | |
532 | for _ in 0..self.mb_w * self.mb_h { | |
533 | self.inter_mbs.push(InterMB::new()); | |
534 | } | |
535 | } | |
536 | if self.golden_mbs.is_empty() { | |
537 | for _ in 0..self.mb_w * self.mb_h { | |
538 | self.golden_mbs.push(InterMB::new()); | |
539 | } | |
540 | } | |
541 | ||
542 | let mut cur_blk = [[0u8; 64]; 6]; | |
543 | ||
544 | let mut mv_est = MVEstimator::new(ref_frame.clone(), mc_buf.clone(), loop_thr, self.me_range); | |
545 | let mut mv_est_g = if let Some(gold_frm) = gold_frame { | |
546 | Some(MVEstimator::new(gold_frm, mc_buf.clone(), loop_thr, self.me_range)) | |
547 | } else { | |
548 | None | |
549 | }; | |
550 | ||
551 | let mut mv_search = self.me_mode.create_search(); | |
552 | ||
553 | let mut tmp_mb = ResidueMB::new(); | |
554 | ||
555 | let mut mb_idx = 0; | |
556 | for mb_y in 0..self.mb_h { | |
557 | for mb_x in 0..self.mb_w { | |
558 | let smb = &self.src_mbs[mb_idx]; | |
559 | ||
560 | let inter_mb = &mut self.inter_mbs[mb_idx]; | |
561 | Self::motion_est_mb(smb, &mut cur_blk, inter_mb, &mut mv_search, &mut mv_est, mb_x, mb_y); | |
562 | inter_mb.residue.fdct(); | |
563 | inter_mb.residue.quant(self.quant); | |
564 | self.mb_types[mb_idx] = VPMBType::InterMV; | |
565 | ||
566 | tmp_mb.dequant_from(&inter_mb.residue, self.quant); | |
567 | tmp_mb.idct(); | |
568 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(inter_mb.reference.iter()) { | |
569 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
570 | *coef = (*coef + add).max(0).min(255); | |
571 | } | |
572 | } | |
573 | let mut best_dist = calc_mb_dist(smb, &tmp_mb); | |
574 | let mut inter_nits = estimate_inter_mb_nits(inter_mb, self.quant, false); | |
575 | if inter_mb.mv[3] != ZERO_MV { | |
576 | inter_nits += estimate_mv_nits(inter_mb.mv[3]); | |
577 | } | |
578 | let mut best_cost = (best_dist as f32) + lambda * (inter_nits as f32); | |
579 | if best_dist > 512 { | |
580 | self.estimate_fourmv(ref_frame.clone(), mc_buf.clone(), mb_idx % self.mb_w, mb_idx / self.mb_w); | |
581 | self.fourmv_mbs[mb_idx].residue.fdct(); | |
582 | self.fourmv_mbs[mb_idx].residue.quant(self.quant); | |
583 | ||
584 | let smb = &self.src_mbs[mb_idx]; | |
585 | tmp_mb.dequant_from(&self.fourmv_mbs[mb_idx].residue, self.quant); | |
586 | tmp_mb.idct(); | |
587 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.fourmv_mbs[mb_idx].reference.iter()) { | |
588 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
589 | *coef = (*coef + add).max(0).min(255); | |
590 | } | |
591 | } | |
592 | let fourmv_dist = calc_mb_dist(smb, &tmp_mb); | |
593 | let fourmv_nits = estimate_inter_mb_nits(&self.fourmv_mbs[mb_idx], self.quant, true); | |
594 | let fourmv_cost = (fourmv_dist as f32) + lambda * (fourmv_nits as f32); | |
595 | if fourmv_cost < best_cost { | |
596 | self.mb_types[mb_idx] = VPMBType::InterFourMV; | |
597 | best_cost = fourmv_cost; | |
598 | best_dist = fourmv_dist; | |
599 | } | |
600 | } | |
601 | let smb = &self.src_mbs[mb_idx]; | |
602 | if best_dist > 512 { | |
603 | if let Some(ref mut mve_gold) = mv_est_g { | |
604 | let gold_mb = &mut self.golden_mbs[mb_idx]; | |
605 | Self::motion_est_mb(smb, &mut cur_blk, gold_mb, &mut mv_search, mve_gold, mb_x, mb_y); | |
606 | gold_mb.residue.fdct(); | |
607 | gold_mb.residue.quant(self.quant); | |
608 | ||
609 | tmp_mb.dequant_from(&gold_mb.residue, self.quant); | |
610 | tmp_mb.idct(); | |
611 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(gold_mb.reference.iter()) { | |
612 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
613 | *coef = (*coef + add).max(0).min(255); | |
614 | } | |
615 | } | |
616 | let golden_dist = calc_mb_dist(smb, &tmp_mb); | |
617 | let golden_nits = estimate_inter_mb_nits(gold_mb, self.quant, false); | |
618 | let golden_cost = (golden_dist as f32) + lambda * (golden_nits as f32); | |
619 | if golden_cost < best_cost { | |
620 | self.mb_types[mb_idx] = VPMBType::GoldenMV; | |
621 | best_cost = golden_cost; | |
622 | best_dist = golden_dist; | |
623 | } | |
624 | } | |
625 | } | |
626 | if best_dist > 512 { | |
627 | let intra_mb = &mut self.intra_mbs[mb_idx]; | |
628 | *intra_mb = smb.clone(); | |
629 | intra_mb.fdct(); | |
630 | for blk in intra_mb.coeffs.iter_mut() { | |
631 | blk[0] -= 4096; | |
632 | } | |
633 | intra_mb.quant(self.quant); | |
634 | ||
635 | tmp_mb.dequant_from(intra_mb, self.quant); | |
636 | tmp_mb.idct(); | |
637 | for blk in tmp_mb.coeffs.iter_mut() { | |
638 | for coef in blk.iter_mut() { | |
639 | *coef = (*coef + 128).max(0).min(255); | |
640 | } | |
641 | } | |
642 | let intra_dist = calc_mb_dist(smb, &tmp_mb); | |
643 | let intra_nits = estimate_intra_mb_nits(&intra_mb.coeffs, self.quant); | |
644 | let intra_cost = (intra_dist as f32) + lambda * (intra_nits as f32); | |
645 | if intra_cost < best_cost { | |
646 | self.mb_types[mb_idx] = VPMBType::Intra; | |
647 | } | |
648 | } | |
649 | ||
650 | mb_idx += 1; | |
651 | } | |
652 | } | |
653 | } | |
654 | pub fn decide_frame_type(&self) -> (bool, bool) { | |
655 | let mut intra_count = 0usize; | |
656 | let mut non_intra = 0usize; | |
657 | for mb_type in self.mb_types.iter() { | |
658 | if mb_type.is_intra() { | |
659 | intra_count += 1; | |
660 | } else { | |
661 | non_intra += 1; | |
662 | } | |
663 | } | |
664 | (intra_count > non_intra * 3, intra_count > non_intra) | |
665 | } | |
666 | fn find_mv_pred(&self, mb_x: usize, mb_y: usize, ref_id: u8) -> (usize, MV, MV, MV) { | |
667 | const CAND_POS: [(i8, i8); 12] = [ | |
668 | (-1, 0), ( 0, -1), | |
669 | (-1, -1), (-1, 1), | |
670 | (-2, 0), ( 0, -2), | |
671 | (-1, -2), (-2, -1), | |
672 | (-2, 1), (-1, 2), | |
673 | (-2, -2), (-2, 2) | |
674 | ]; | |
675 | ||
676 | let mut nearest_mv = ZERO_MV; | |
677 | let mut near_mv = ZERO_MV; | |
678 | let mut pred_mv = ZERO_MV; | |
679 | let mut num_mv: usize = 0; | |
680 | ||
681 | for (i, (yoff, xoff)) in CAND_POS.iter().enumerate() { | |
682 | let cx = (mb_x as isize) + (*xoff as isize); | |
683 | let cy = (mb_y as isize) + (*yoff as isize); | |
684 | if (cx < 0) || (cy < 0) { | |
685 | continue; | |
686 | } | |
687 | let cx = cx as usize; | |
688 | let cy = cy as usize; | |
689 | if (cx >= self.mb_w) || (cy >= self.mb_h) { | |
690 | continue; | |
691 | } | |
692 | let mb_pos = cx + cy * self.mb_w; | |
693 | let mv = match self.mb_types[mb_pos].get_ref_id() { | |
694 | 0 => ZERO_MV, | |
695 | 1 => if self.mb_types[mb_pos] != VPMBType::InterFourMV { | |
696 | self.inter_mbs[mb_pos].mv[3] | |
697 | } else { | |
698 | self.fourmv_mbs[mb_pos].mv[3] | |
699 | }, | |
700 | _ => self.golden_mbs[mb_pos].mv[3], | |
701 | }; | |
702 | if (self.mb_types[mb_pos].get_ref_id() != ref_id) || (mv == ZERO_MV) { | |
703 | continue; | |
704 | } | |
705 | if num_mv == 0 { | |
706 | nearest_mv = mv; | |
707 | num_mv += 1; | |
708 | if i < 2 { | |
709 | pred_mv = mv; | |
710 | } | |
711 | } else if mv != nearest_mv { | |
712 | near_mv = mv; | |
713 | num_mv += 1; | |
714 | break; | |
715 | } | |
716 | } | |
717 | ||
718 | (num_mv, nearest_mv, near_mv, pred_mv) | |
719 | } | |
720 | pub fn predict_mvs(&mut self) { | |
721 | let mut mb_idx = 0; | |
722 | self.num_mv.clear(); | |
723 | if self.coded_mv.is_empty() { | |
724 | self.coded_mv.resize(self.mb_w * self.mb_h, [ZERO_MV; 4]); | |
725 | } | |
726 | for mb_y in 0..self.mb_h { | |
727 | for mb_x in 0..self.mb_w { | |
728 | let (num_mv, nearest_mv, near_mv, pred_mv) = self.find_mv_pred(mb_x, mb_y, VP_REF_INTER); | |
729 | let mb_type = self.mb_types[mb_idx]; | |
730 | self.num_mv.push(num_mv as u8); | |
731 | let golden = mb_type.get_ref_id() == VP_REF_GOLDEN; | |
732 | let mv = if !golden { self.inter_mbs[mb_idx].mv[3] } else { self.golden_mbs[mb_idx].mv[3] }; | |
733 | ||
734 | let mb_type = if mb_type == VPMBType::Intra { | |
735 | VPMBType::Intra | |
736 | } else if mb_type == VPMBType::InterFourMV { | |
737 | for i in 0..4 { | |
738 | let mv = self.fourmv_mbs[mb_idx].mv[i]; | |
739 | self.coded_mv[mb_idx][i] = ZERO_MV; | |
740 | if mv == ZERO_MV { | |
741 | self.fmv_sub[mb_idx][i] = VPMBType::InterNoMV; | |
742 | } else { | |
743 | self.fmv_sub[mb_idx][i] = match num_mv { | |
744 | 0 => { | |
745 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
746 | VPMBType::InterMV | |
747 | }, | |
748 | 1 => { | |
749 | if nearest_mv == mv { | |
750 | VPMBType::InterNearest | |
751 | } else { | |
752 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
753 | VPMBType::InterMV | |
754 | } | |
755 | }, | |
756 | _ => { | |
757 | if nearest_mv == mv { | |
758 | VPMBType::InterNearest | |
759 | } else if near_mv == mv { | |
760 | VPMBType::InterNear | |
761 | } else { | |
762 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
763 | VPMBType::InterMV | |
764 | } | |
765 | }, | |
766 | }; | |
767 | } | |
768 | } | |
769 | VPMBType::InterFourMV | |
770 | } else if mv == ZERO_MV { | |
771 | if !golden { | |
772 | VPMBType::InterNoMV | |
773 | } else { | |
774 | VPMBType::GoldenNoMV | |
775 | } | |
776 | } else if mb_type.get_ref_id() == VP_REF_INTER { | |
777 | self.coded_mv[mb_idx][3] = mv; | |
778 | match num_mv { | |
779 | 0 => VPMBType::InterMV, | |
780 | 1 => { | |
781 | if nearest_mv == mv { | |
782 | VPMBType::InterNearest | |
783 | } else { | |
784 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
785 | VPMBType::InterMV | |
786 | } | |
787 | }, | |
788 | _ => { | |
789 | if nearest_mv == mv { | |
790 | VPMBType::InterNearest | |
791 | } else if near_mv == mv { | |
792 | VPMBType::InterNear | |
793 | } else { | |
794 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
795 | VPMBType::InterMV | |
796 | } | |
797 | }, | |
798 | } | |
799 | } else { | |
800 | let (num_mv, nearest_mv, near_mv, pred_mv) = self.find_mv_pred(mb_x, mb_y, VP_REF_GOLDEN); | |
801 | self.coded_mv[mb_idx][3] = ZERO_MV; | |
802 | match num_mv { | |
803 | 0 => { | |
804 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
805 | VPMBType::GoldenMV | |
806 | }, | |
807 | 1 => { | |
808 | if nearest_mv == mv { | |
809 | VPMBType::GoldenNearest | |
810 | } else { | |
811 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
812 | VPMBType::GoldenMV | |
813 | } | |
814 | }, | |
815 | _ => { | |
816 | if nearest_mv == mv { | |
817 | VPMBType::GoldenNearest | |
818 | } else if near_mv == mv { | |
819 | VPMBType::GoldenNear | |
820 | } else { | |
821 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
822 | VPMBType::GoldenMV | |
823 | } | |
824 | }, | |
825 | } | |
826 | }; | |
827 | self.mb_types[mb_idx] = mb_type; | |
828 | mb_idx += 1; | |
829 | } | |
830 | } | |
831 | } | |
832 | pub fn apply_dc_prediction(&mut self, dc_pred: &mut VP56DCPred) { | |
833 | dc_pred.reset(); | |
834 | ||
835 | let mut mb_idx = 0; | |
836 | for _mb_y in 0..self.mb_h { | |
837 | for _mb_x in 0..self.mb_w { | |
838 | let mb_type = self.mb_types[mb_idx]; | |
839 | let mb = self.get_mb_mut(mb_idx); | |
840 | for (i, blk) in mb.coeffs.iter_mut().enumerate() { | |
841 | dc_pred.predict_dc(mb_type, i, blk, true); | |
842 | } | |
843 | dc_pred.next_mb(); | |
844 | mb_idx += 1; | |
845 | } | |
846 | dc_pred.update_row(); | |
847 | } | |
848 | } | |
849 | } |