]>
Commit | Line | Data |
---|---|---|
3952bfd9 KS |
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: Box<dyn MVSearch> = match self.me_mode { | |
359 | MVSearchMode::Full => Box::new(FullMVSearch::new()), | |
360 | MVSearchMode::Diamond => Box::new(DiaSearch::new()), | |
361 | MVSearchMode::Hexagon => Box::new(HexSearch::new()), | |
362 | }; | |
363 | let mut mb_pos = 0; | |
364 | for (mb_y, row) in inter_mbs.chunks_mut(self.mb_w).enumerate() { | |
365 | for (mb_x, mb) in row.iter_mut().enumerate() { | |
366 | self.src_mbs[mb_pos].fill(&mut cur_blk); | |
367 | ||
368 | let (best_mv, _best_dist) = mv_search.search_mb(&mut mv_est, &cur_blk, mb_x, mb_y); | |
369 | mb.mv[3] = best_mv; | |
370 | ||
371 | for i in 0..4 { | |
372 | mv_est.mc_block(i, 0, mb_x * 16 + (i & 1) * 8, mb_y * 16 + (i >> 1) * 8, best_mv); | |
373 | sub_blk(&mut mb.residue.coeffs[i], &cur_blk[i], &mv_est.ref_blk[i]); | |
374 | } | |
375 | for plane in 1..3 { | |
376 | mv_est.mc_block(plane + 3, plane, mb_x * 8, mb_y * 8, best_mv); | |
377 | sub_blk(&mut mb.residue.coeffs[plane + 3], &cur_blk[plane + 3], &mv_est.ref_blk[plane + 3]); | |
378 | } | |
379 | ||
380 | for (dblk, sblk) in mb.reference.iter_mut().zip(mv_est.ref_blk.iter()) { | |
381 | for (dst, &src) in dblk.iter_mut().zip(sblk.iter()) { | |
382 | *dst = i16::from(src); | |
383 | } | |
384 | } | |
385 | mb_pos += 1; | |
386 | } | |
387 | } | |
388 | } | |
389 | fn estimate_fourmv(&mut self, ref_frame: NAVideoBufferRef<u8>, mc_buf: NAVideoBufferRef<u8>, mb_x: usize, mb_y: usize) -> bool { | |
390 | let loop_thr = i16::from(VP56_FILTER_LIMITS[self.quant as usize]); | |
391 | ||
392 | if self.fourmv_mbs.is_empty() { | |
393 | for _ in 0..self.mb_w * self.mb_h { | |
394 | self.fourmv_mbs.push(InterMB::new()); | |
395 | } | |
396 | } | |
397 | if self.fmv_sub.is_empty() { | |
398 | self.fmv_sub.resize(self.mb_w * self.mb_h, [VPMBType::Intra; 4]); | |
399 | } | |
400 | ||
401 | let mb_pos = mb_x + mb_y * self.mb_w; | |
402 | let mb = &mut self.fourmv_mbs[mb_pos]; | |
403 | ||
404 | let mut cur_blk = [[0u8; 64]; 6]; | |
405 | self.src_mbs[mb_pos].fill(&mut cur_blk); | |
406 | ||
407 | let mut mv_est = MVEstimator::new(ref_frame, mc_buf, loop_thr, self.me_range); | |
408 | ||
409 | let mut mv_search: Box<dyn MVSearch> = match self.me_mode { | |
410 | MVSearchMode::Full => Box::new(FullMVSearch::new()), | |
411 | MVSearchMode::Diamond => Box::new(DiaSearch::new()), | |
412 | MVSearchMode::Hexagon => Box::new(HexSearch::new()), | |
413 | }; | |
414 | ||
415 | for i in 0..4 { | |
416 | let xpos = mb_x * 16 + (i & 1) * 8; | |
417 | let ypos = mb_y * 16 + (i >> 1) * 8; | |
418 | let (best_mv, _best_dist) = mv_search.search_blk(&mut mv_est, &cur_blk[i], xpos, ypos); | |
419 | mb.mv[i] = best_mv; | |
420 | } | |
421 | let mvsum = mb.mv[0] + mb.mv[1] + mb.mv[2] + mb.mv[3]; | |
422 | let chroma_mv = MV{ x: mvsum.x / 4, y: mvsum.y / 4}; | |
423 | ||
424 | for (i, blk) in mb.residue.coeffs[..4].iter_mut().enumerate() { | |
425 | let xpos = mb_x * 16 + (i & 1) * 8; | |
426 | let ypos = mb_y * 16 + (i >> 1) * 8; | |
427 | mv_est.mc_block(i, 0, xpos, ypos, mb.mv[i]); | |
428 | sub_blk(blk, &cur_blk[i], &mv_est.ref_blk[i]); | |
429 | } | |
430 | for plane in 1..3 { | |
431 | mv_est.mc_block(plane + 3, plane, mb_x * 8, mb_y * 8, chroma_mv); | |
432 | sub_blk(&mut mb.residue.coeffs[plane + 3], &cur_blk[plane + 3], &mv_est.ref_blk[plane + 3]); | |
433 | } | |
434 | ||
435 | for (dblk, sblk) in mb.reference.iter_mut().zip(mv_est.ref_blk.iter()) { | |
436 | for (dst, &src) in dblk.iter_mut().zip(sblk.iter()) { | |
437 | *dst = i16::from(src); | |
438 | } | |
439 | } | |
440 | ||
441 | (mb.mv[0] != mb.mv[1]) || (mb.mv[0] != mb.mv[2]) || (mb.mv[0] != mb.mv[3]) | |
442 | } | |
443 | pub fn select_inter_blocks(&mut self, ref_frame: NAVideoBufferRef<u8>, mc_buf: NAVideoBufferRef<u8>, has_golden_frame: bool, lambda: f32) { | |
444 | let mut tmp_mb = ResidueMB::new(); | |
445 | for mb_idx in 0..self.mb_w * self.mb_h { | |
446 | tmp_mb.dequant_from(&self.intra_mbs[mb_idx], self.quant); | |
447 | tmp_mb.idct(); | |
448 | for blk in tmp_mb.coeffs.iter_mut() { | |
449 | for coef in blk.iter_mut() { | |
450 | *coef = (*coef + 128).max(0).min(255); | |
451 | } | |
452 | } | |
453 | let intra_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
454 | let intra_nits = estimate_intra_mb_nits(&self.intra_mbs[mb_idx].coeffs, self.quant); | |
455 | let intra_cost = (intra_dist as f32) + lambda * (intra_nits as f32); | |
456 | ||
457 | tmp_mb.dequant_from(&self.inter_mbs[mb_idx].residue, self.quant); | |
458 | tmp_mb.idct(); | |
459 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.inter_mbs[mb_idx].reference.iter()) { | |
460 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
461 | *coef = (*coef + add).max(0).min(255); | |
462 | } | |
463 | } | |
464 | let inter_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
465 | let mut inter_nits = estimate_inter_mb_nits(&self.inter_mbs[mb_idx], self.quant, false); | |
466 | if self.inter_mbs[mb_idx].mv[3] != ZERO_MV { | |
467 | inter_nits += estimate_mv_nits(self.inter_mbs[mb_idx].mv[3]); | |
468 | } | |
469 | let mut inter_cost = (inter_dist as f32) + lambda * (inter_nits as f32); | |
470 | ||
471 | if inter_cost < intra_cost { | |
472 | self.mb_types[mb_idx] = VPMBType::InterMV; | |
473 | ||
474 | if inter_dist > 512 { | |
475 | self.estimate_fourmv(ref_frame.clone(), mc_buf.clone(), mb_idx % self.mb_w, mb_idx / self.mb_w); | |
476 | self.fourmv_mbs[mb_idx].residue.fdct(); | |
477 | self.fourmv_mbs[mb_idx].residue.quant(self.quant); | |
478 | ||
479 | tmp_mb.dequant_from(&self.fourmv_mbs[mb_idx].residue, self.quant); | |
480 | tmp_mb.idct(); | |
481 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.fourmv_mbs[mb_idx].reference.iter()) { | |
482 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
483 | *coef = (*coef + add).max(0).min(255); | |
484 | } | |
485 | } | |
486 | let fourmv_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
487 | let fourmv_nits = estimate_inter_mb_nits(&self.fourmv_mbs[mb_idx], self.quant, true); | |
488 | let fourmv_cost = (fourmv_dist as f32) + lambda * (fourmv_nits as f32); | |
489 | if fourmv_cost < inter_cost { | |
490 | self.mb_types[mb_idx] = VPMBType::InterFourMV; | |
491 | inter_cost = fourmv_cost; | |
492 | } | |
493 | } | |
494 | } | |
495 | ||
496 | if has_golden_frame { | |
497 | tmp_mb.dequant_from(&self.golden_mbs[mb_idx].residue, self.quant); | |
498 | tmp_mb.idct(); | |
499 | for (blk, res) in tmp_mb.coeffs.iter_mut().zip(self.golden_mbs[mb_idx].reference.iter()) { | |
500 | for (coef, add) in blk.iter_mut().zip(res.iter()) { | |
501 | *coef = (*coef + add).max(0).min(255); | |
502 | } | |
503 | } | |
504 | let golden_dist = calc_mb_dist(&self.src_mbs[mb_idx], &tmp_mb); | |
505 | let golden_nits = estimate_inter_mb_nits(&self.golden_mbs[mb_idx], self.quant, false); | |
506 | let golden_cost = (golden_dist as f32) + lambda * (golden_nits as f32); | |
507 | ||
508 | if (self.mb_types[mb_idx].is_intra() && golden_cost < intra_cost) || | |
509 | (!self.mb_types[mb_idx].is_intra() && golden_cost < inter_cost) { | |
510 | self.mb_types[mb_idx] = VPMBType::GoldenMV; | |
511 | } | |
512 | } | |
513 | } | |
514 | } | |
515 | pub fn decide_frame_type(&self) -> (bool, bool) { | |
516 | let mut intra_count = 0usize; | |
517 | let mut non_intra = 0usize; | |
518 | for mb_type in self.mb_types.iter() { | |
519 | if mb_type.is_intra() { | |
520 | intra_count += 1; | |
521 | } else { | |
522 | non_intra += 1; | |
523 | } | |
524 | } | |
525 | (intra_count > non_intra * 3, intra_count > non_intra) | |
526 | } | |
527 | fn find_mv_pred(&self, mb_x: usize, mb_y: usize, ref_id: u8) -> (usize, MV, MV, MV) { | |
528 | const CAND_POS: [(i8, i8); 12] = [ | |
529 | (-1, 0), ( 0, -1), | |
530 | (-1, -1), (-1, 1), | |
531 | (-2, 0), ( 0, -2), | |
532 | (-1, -2), (-2, -1), | |
533 | (-2, 1), (-1, 2), | |
534 | (-2, -2), (-2, 2) | |
535 | ]; | |
536 | ||
537 | let mut nearest_mv = ZERO_MV; | |
538 | let mut near_mv = ZERO_MV; | |
539 | let mut pred_mv = ZERO_MV; | |
540 | let mut num_mv: usize = 0; | |
541 | ||
542 | for (i, (yoff, xoff)) in CAND_POS.iter().enumerate() { | |
543 | let cx = (mb_x as isize) + (*xoff as isize); | |
544 | let cy = (mb_y as isize) + (*yoff as isize); | |
545 | if (cx < 0) || (cy < 0) { | |
546 | continue; | |
547 | } | |
548 | let cx = cx as usize; | |
549 | let cy = cy as usize; | |
550 | if (cx >= self.mb_w) || (cy >= self.mb_h) { | |
551 | continue; | |
552 | } | |
553 | let mb_pos = cx + cy * self.mb_w; | |
554 | let mv = match self.mb_types[mb_pos].get_ref_id() { | |
555 | 0 => ZERO_MV, | |
556 | 1 => if self.mb_types[mb_pos] != VPMBType::InterFourMV { | |
557 | self.inter_mbs[mb_pos].mv[3] | |
558 | } else { | |
559 | self.fourmv_mbs[mb_pos].mv[3] | |
560 | }, | |
561 | _ => self.golden_mbs[mb_pos].mv[3], | |
562 | }; | |
563 | if (self.mb_types[mb_pos].get_ref_id() != ref_id) || (mv == ZERO_MV) { | |
564 | continue; | |
565 | } | |
566 | if num_mv == 0 { | |
567 | nearest_mv = mv; | |
568 | num_mv += 1; | |
569 | if i < 2 { | |
570 | pred_mv = mv; | |
571 | } | |
572 | } else if mv != nearest_mv { | |
573 | near_mv = mv; | |
574 | num_mv += 1; | |
575 | break; | |
576 | } | |
577 | } | |
578 | ||
579 | (num_mv, nearest_mv, near_mv, pred_mv) | |
580 | } | |
581 | pub fn predict_mvs(&mut self) { | |
582 | let mut mb_idx = 0; | |
583 | self.num_mv.clear(); | |
584 | if self.coded_mv.is_empty() { | |
585 | self.coded_mv.resize(self.mb_w * self.mb_h, [ZERO_MV; 4]); | |
586 | } | |
587 | for mb_y in 0..self.mb_h { | |
588 | for mb_x in 0..self.mb_w { | |
589 | let (num_mv, nearest_mv, near_mv, pred_mv) = self.find_mv_pred(mb_x, mb_y, VP_REF_INTER); | |
590 | let mb_type = self.mb_types[mb_idx]; | |
591 | self.num_mv.push(num_mv as u8); | |
592 | let golden = mb_type.get_ref_id() == VP_REF_GOLDEN; | |
593 | let mv = if !golden { self.inter_mbs[mb_idx].mv[3] } else { self.golden_mbs[mb_idx].mv[3] }; | |
594 | ||
595 | let mb_type = if mb_type == VPMBType::Intra { | |
596 | VPMBType::Intra | |
597 | } else if mb_type == VPMBType::InterFourMV { | |
598 | for i in 0..4 { | |
599 | let mv = self.fourmv_mbs[mb_idx].mv[i]; | |
600 | self.coded_mv[mb_idx][i] = ZERO_MV; | |
601 | if mv == ZERO_MV { | |
602 | self.fmv_sub[mb_idx][i] = VPMBType::InterNoMV; | |
603 | } else { | |
604 | self.fmv_sub[mb_idx][i] = match num_mv { | |
605 | 0 => { | |
606 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
607 | VPMBType::InterMV | |
608 | }, | |
609 | 1 => { | |
610 | if nearest_mv == mv { | |
611 | VPMBType::InterNearest | |
612 | } else { | |
613 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
614 | VPMBType::InterMV | |
615 | } | |
616 | }, | |
617 | _ => { | |
618 | if nearest_mv == mv { | |
619 | VPMBType::InterNearest | |
620 | } else if near_mv == mv { | |
621 | VPMBType::InterNear | |
622 | } else { | |
623 | self.coded_mv[mb_idx][i] = mv - pred_mv; | |
624 | VPMBType::InterMV | |
625 | } | |
626 | }, | |
627 | }; | |
628 | } | |
629 | } | |
630 | VPMBType::InterFourMV | |
631 | } else if mv == ZERO_MV { | |
632 | if !golden { | |
633 | VPMBType::InterNoMV | |
634 | } else { | |
635 | VPMBType::GoldenNoMV | |
636 | } | |
637 | } else if mb_type.get_ref_id() == VP_REF_INTER { | |
638 | self.coded_mv[mb_idx][3] = mv; | |
639 | match num_mv { | |
640 | 0 => VPMBType::InterMV, | |
641 | 1 => { | |
642 | if nearest_mv == mv { | |
643 | VPMBType::InterNearest | |
644 | } else { | |
645 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
646 | VPMBType::InterMV | |
647 | } | |
648 | }, | |
649 | _ => { | |
650 | if nearest_mv == mv { | |
651 | VPMBType::InterNearest | |
652 | } else if near_mv == mv { | |
653 | VPMBType::InterNear | |
654 | } else { | |
655 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
656 | VPMBType::InterMV | |
657 | } | |
658 | }, | |
659 | } | |
660 | } else { | |
661 | let (num_mv, nearest_mv, near_mv, pred_mv) = self.find_mv_pred(mb_x, mb_y, VP_REF_GOLDEN); | |
662 | self.coded_mv[mb_idx][3] = ZERO_MV; | |
663 | match num_mv { | |
664 | 0 => { | |
665 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
666 | VPMBType::GoldenMV | |
667 | }, | |
668 | 1 => { | |
669 | if nearest_mv == mv { | |
670 | VPMBType::GoldenNearest | |
671 | } else { | |
672 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
673 | VPMBType::GoldenMV | |
674 | } | |
675 | }, | |
676 | _ => { | |
677 | if nearest_mv == mv { | |
678 | VPMBType::GoldenNearest | |
679 | } else if near_mv == mv { | |
680 | VPMBType::GoldenNear | |
681 | } else { | |
682 | self.coded_mv[mb_idx][3] = mv - pred_mv; | |
683 | VPMBType::GoldenMV | |
684 | } | |
685 | }, | |
686 | } | |
687 | }; | |
688 | self.mb_types[mb_idx] = mb_type; | |
689 | mb_idx += 1; | |
690 | } | |
691 | } | |
692 | } | |
693 | pub fn apply_dc_prediction(&mut self, dc_pred: &mut VP56DCPred) { | |
694 | dc_pred.reset(); | |
695 | ||
696 | let mut mb_idx = 0; | |
697 | for _mb_y in 0..self.mb_h { | |
698 | for _mb_x in 0..self.mb_w { | |
699 | let mb_type = self.mb_types[mb_idx]; | |
700 | let mb = self.get_mb_mut(mb_idx); | |
701 | for (i, blk) in mb.coeffs.iter_mut().enumerate() { | |
702 | dc_pred.predict_dc(mb_type, i, blk, true); | |
703 | } | |
704 | dc_pred.next_mb(); | |
705 | mb_idx += 1; | |
706 | } | |
707 | dc_pred.update_row(); | |
708 | } | |
709 | } | |
710 | } |