1 use nihav_core::frame::*;
2 use nihav_codec_support::codecs::{MV, ZERO_MV};
4 use super::dsp::{RefMBData, luma_mc, chroma_mc};
6 #[derive(Clone,Copy,PartialEq,Default)]
7 pub enum MVSearchMode {
16 pub const fn get_possible_modes() -> &'static [&'static str] {
17 &["diamond", "hexagon", "umh"]
19 fn create(self) -> Box<dyn MVSearch+Send> {
21 MVSearchMode::Dummy => Box::new(DummySearcher{}),
22 MVSearchMode::Diamond => Box::new(DiaSearch::new()),
23 MVSearchMode::Hexagon => Box::new(HexSearch::new()),
24 MVSearchMode::UMH => Box::new(UnevenHexSearch::new()),
29 impl std::fmt::Display for MVSearchMode {
30 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
32 MVSearchMode::Diamond => write!(f, "diamond"),
33 MVSearchMode::Hexagon => write!(f, "hexagon"),
34 MVSearchMode::UMH => write!(f, "umh"),
35 MVSearchMode::Dummy => write!(f, "dummy"),
40 impl FromStr for MVSearchMode {
42 fn from_str(s: &str) -> Result<Self, Self::Err> {
44 "diamond" => Ok(MVSearchMode::Diamond),
45 "hexagon" => Ok(MVSearchMode::Hexagon),
46 "umh" => Ok(MVSearchMode::UMH),
47 "dummy" => Ok(MVSearchMode::Dummy),
53 const MAX_DIST: u32 = std::u32::MAX;
54 const DIST_THRESH: u32 = 256;
57 fn from_pixels(self) -> Self;
60 impl FromPixels for MV {
61 fn from_pixels(self) -> MV {
62 MV { x: self.x * 4, y: self.y * 4 }
66 const DIA_PATTERN: [MV; 9] = [
78 const HEX_PATTERN: [MV; 7] = [
88 const REFINEMENT: [MV; 4] = [
95 macro_rules! search_template {
96 ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident, $threshold: expr) => ({
97 search_template!($self, $mv_est, $cur_blk, $mb_x, $mb_y, $sad_func, $threshold, ZERO_MV, MAX_DIST, true)
99 ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident, $threshold: expr, $start_mv: expr, $best_dist: expr, $fullpel_stage: expr) => ({
100 let mut best_dist = $best_dist;
101 let mut best_mv = $start_mv;
109 let mut cur_best_dist = best_dist;
110 for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
111 if *dist == MAX_DIST {
112 *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point.from_pixels(), cur_best_dist);
113 cur_best_dist = cur_best_dist.min(*dist);
114 if *dist <= $threshold {
119 min_dist = $self.dist[0];
121 for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
125 if dist <= $threshold {
130 if min_dist <= $threshold || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range || $self.point[min_idx].y.abs() >= $mv_est.mv_range {
133 best_dist = min_dist;
134 $self.update($self.steps[min_idx]);
136 best_dist = min_dist;
137 best_mv = $self.point[min_idx];
138 if best_dist <= $threshold {
139 return (best_mv.from_pixels(), best_dist);
141 for &step in REFINEMENT.iter() {
142 let mv = best_mv + step;
143 let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv.from_pixels(), MAX_DIST);
144 if best_dist > dist {
149 best_mv = best_mv.from_pixels();
150 if best_dist <= $threshold {
151 return (best_mv, best_dist);
156 $self.set_new_point(best_mv, best_dist);
158 let mut cur_best_dist = best_dist;
159 for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
160 if *dist == MAX_DIST {
161 *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point, cur_best_dist);
162 cur_best_dist = cur_best_dist.min(*dist);
163 if *dist <= $threshold {
168 min_dist = $self.dist[0];
170 for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
174 if dist <= $threshold {
179 if min_dist <= $threshold || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range * 8 || $self.point[min_idx].y.abs() >= $mv_est.mv_range * 8 {
182 best_dist = min_dist;
183 $self.update($self.steps[min_idx]);
185 best_dist = min_dist;
186 best_mv = $self.point[min_idx];
187 if best_dist <= $threshold {
188 return (best_mv, best_dist);
190 for &step in REFINEMENT.iter() {
191 let mv = best_mv + step;
192 let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv, MAX_DIST);
193 if best_dist > dist {
202 macro_rules! pattern_search {
203 ($struct_name: ident, $patterns: expr) => {
204 pub struct $struct_name {
205 point: [MV; $patterns.len()],
206 dist: [u32; $patterns.len()],
207 steps: &'static [MV; $patterns.len()],
211 pub fn new() -> Self {
214 dist: [MAX_DIST; $patterns.len()],
218 fn reset(&mut self) {
219 self.point = $patterns;
220 self.dist = [MAX_DIST; $patterns.len()];
222 fn set_new_point(&mut self, start: MV, dist: u32) {
223 for (dst, &src) in self.point.iter_mut().zip(self.steps.iter()) {
226 self.dist = [MAX_DIST; $patterns.len()];
229 fn update(&mut self, step: MV) {
230 let mut new_point = self.point;
231 let mut new_dist = [MAX_DIST; $patterns.len()];
233 for point in new_point.iter_mut() {
237 for (new_point, new_dist) in new_point.iter_mut().zip(new_dist.iter_mut()) {
238 for (&old_point, &old_dist) in self.point.iter().zip(self.dist.iter()) {
239 if *new_point == old_point {
240 *new_dist = old_dist;
245 self.point = new_point;
246 self.dist = new_dist;
250 impl MVSearch for $struct_name {
251 fn search_mb(&mut self, mv_est: &mut MVEstimator, cur_mb: &RefMBData, mb_x: usize, mb_y: usize, _cand_mvs: &[MV]) -> (MV, u32) {
252 search_template!(self, mv_est, cur_mb, mb_x, mb_y, sad_mb, DIST_THRESH)
254 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, _cand_mvs: &[MV]) -> (MV, u32) {
255 search_template!(self, mv_est, ref_blk, xpos, ypos, sad_blk8, DIST_THRESH / 4)
261 pattern_search!(DiaSearch, DIA_PATTERN);
262 pattern_search!(HexSearch, HEX_PATTERN);
264 const LARGE_HEX_PATTERN: [MV; 16] = [
283 const UNSYMM_CROSS: [MV; 4] = [
291 struct UniqueSet<T:Copy+Default> {
296 impl<T:Copy+Default+PartialEq> UniqueSet<T> {
297 fn new() -> Self { Self::default() }
298 fn clear(&mut self) { self.count = 0; }
299 fn get_list(&self) -> &[T] { &self.list[..self.count] }
300 fn add(&mut self, val: T) {
301 if self.count < self.list.len() && !self.get_list().contains(&val) {
302 self.list[self.count] = val;
309 fn scale(self, scale: i16) -> Self;
310 fn is_in_range(self, range: i16) -> bool;
314 fn scale(self, scale: i16) -> MV {
315 MV { x: self.x * scale, y: self.y * scale }
317 fn is_in_range(self, range: i16) -> bool {
318 self.x.abs() <= range && self.y.abs() <= range
322 macro_rules! single_search_step {
323 ($start:expr, $best_dist:expr, $mv_est:expr, $sad_func:ident, $ref_blk:expr, $xpos:expr, $ypos:expr, $pattern:expr, $scale:expr, $dist_thr:expr) => {{
324 let mut best_mv = $start;
325 let mut best_dist = $best_dist;
326 for point in $pattern.iter() {
327 let mv = point.scale($scale) + $start;
328 if !mv.is_in_range($mv_est.mv_range * 4) {
331 let dist = $mv_est.$sad_func($ref_blk, $xpos, $ypos, mv, best_dist);
332 if dist < best_dist {
335 if best_dist < $dist_thr {
340 (best_mv, best_dist, best_mv != $start)
344 struct UnevenHexSearch {
345 mv_list: UniqueSet<MV>,
348 impl UnevenHexSearch {
351 mv_list: UniqueSet::new(),
354 fn get_cand_mv(&mut self, cand_mvs: &[MV]) -> MV {
355 self.mv_list.clear();
356 for &mv in cand_mvs.iter() {
357 self.mv_list.add(mv);
359 match self.mv_list.count {
360 1 => self.mv_list.list[0],
361 3 => MV::pred(self.mv_list.list[0], self.mv_list.list[1], self.mv_list.list[2]),
363 let sum = self.mv_list.get_list().iter().fold((0i32, 0i32),
364 |acc, mv| (acc.0 + i32::from(mv.x), acc.1 + i32::from(mv.y)));
365 MV {x: (sum.0 / (self.mv_list.count as i32)) as i16,
366 y: (sum.1 / (self.mv_list.count as i32)) as i16}
372 macro_rules! umh_search_template {
373 ($cand_mv:expr, $cutoff:expr, $mv_est:expr, $sad_func:ident, $ref_blk:expr, $xpos:expr, $ypos:expr) => {{
374 let cand_mv = $cand_mv;
375 let best_dist = $mv_est.$sad_func($ref_blk, $xpos, $ypos, cand_mv, MAX_DIST);
376 if best_dist < $cutoff {
377 return (cand_mv, best_dist);
380 // step 1 - small refinement search
381 let (mut cand_mv, mut best_dist, _) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, DIA_PATTERN, 1, $cutoff);
382 if best_dist < $cutoff {
383 return (cand_mv, best_dist);
386 // step 2 - unsymmetrical cross search
388 let (mv, dist, changed) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, UNSYMM_CROSS, 4, $cutoff);
394 if best_dist < $cutoff {
399 // step 3 - multi-hexagon grid search
402 let (mv, dist, changed) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, LARGE_HEX_PATTERN, scale, $cutoff);
408 if best_dist < $cutoff {
413 // step 4 - final hexagon search
414 let (cand_mv, best_dist, _) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, HEX_PATTERN, 1, $cutoff);
415 if best_dist > $cutoff {
416 let (mv, dist, _) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, DIA_PATTERN, 1, $cutoff);
424 impl MVSearch for UnevenHexSearch {
425 fn search_mb(&mut self, mv_est: &mut MVEstimator, cur_mb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32) {
426 let cand_mv = self.get_cand_mv(cand_mvs);
427 let cutoff = mv_est.cutoff_thr;
428 umh_search_template!(cand_mv, cutoff, mv_est, sad_mb, cur_mb, mb_x, mb_y)
430 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, cand_mvs: &[MV]) -> (MV, u32) {
431 let cand_mv = self.get_cand_mv(cand_mvs);
432 let cutoff = mv_est.cutoff_thr / 4;
433 umh_search_template!(cand_mv, cutoff, mv_est, sad_blk8, ref_blk, xpos, ypos)
437 struct MVEstimator<'a> {
438 pic: &'a NAVideoBuffer<u8>,
444 ($src1:expr, $src2:expr) => {
445 $src1.iter().zip($src2.iter()).fold(0u32, |acc, (&a, &b)|
446 acc + (((i32::from(a) - i32::from(b)) * (i32::from(a) - i32::from(b))) as u32))
450 impl<'a> MVEstimator<'a> {
451 fn sad_mb(&self, ref_mb: &RefMBData, mb_x: usize, mb_y: usize, mv: MV, cur_best_dist: u32) -> u32 {
452 let mut dst = RefMBData::new();
453 luma_mc(&mut dst.y, 16, self.pic, mb_x * 16, mb_y * 16, mv, true);
456 for (dline, sline) in dst.y.chunks(16).zip(ref_mb.y.chunks(16)) {
457 dist += sad!(dline, sline);
458 if dist > cur_best_dist {
462 chroma_mc(&mut dst.u, 8, self.pic, mb_x * 8, mb_y * 8, 1, mv, true);
463 dist += sad!(dst.u, ref_mb.u);
464 if dist > cur_best_dist {
467 chroma_mc(&mut dst.v, 8, self.pic, mb_x * 8, mb_y * 8, 2, mv, true);
468 dist += sad!(dst.v, ref_mb.v);
472 fn sad_blk8(&self, ref_mb: &RefMBData, xpos: usize, ypos: usize, mv: MV, cur_best_dist: u32) -> u32 {
473 let mut cur_y = [0; 64];
474 let mut cur_u = [0; 16];
475 let mut cur_v = [0; 16];
479 let y_off = (xpos & 8) + (ypos & 8) * 16;
480 luma_mc(&mut cur_y, 8, self.pic, xpos, ypos, mv, false);
481 for (dline, sline) in cur_y.chunks(8).zip(ref_mb.y[y_off..].chunks(16)) {
482 dist += sad!(dline, sline);
483 if dist > cur_best_dist {
488 let c_off = (xpos & 8) / 2 + (ypos & 8) * 4;
489 chroma_mc(&mut cur_u, 4, self.pic, xpos / 2, ypos / 2, 1, mv, false);
490 for (dline, sline) in cur_u.chunks(4).zip(ref_mb.u[c_off..].chunks(8)) {
491 dist += sad!(dline, sline);
492 if dist > cur_best_dist {
496 chroma_mc(&mut cur_v, 4, self.pic, xpos / 2, ypos / 2, 2, mv, false);
497 for (dline, sline) in cur_v.chunks(4).zip(ref_mb.v[c_off..].chunks(8)) {
498 dist += sad!(dline, sline);
499 if dist > cur_best_dist {
509 fn search_mb(&mut self, mv_est: &mut MVEstimator, ref_mb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32);
510 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, cand_mvs: &[MV]) -> (MV, u32);
513 struct DummySearcher {}
515 impl MVSearch for DummySearcher {
516 fn search_mb(&mut self, _mv_est: &mut MVEstimator, _ref_mb: &RefMBData, _mb_x: usize, _mb_y: usize, _cand_mvs: &[MV]) -> (MV, u32) {
517 (ZERO_MV, std::u32::MAX / 2)
519 fn search_blk8(&mut self, _mv_est: &mut MVEstimator, _ref_mb: &RefMBData, _xpos: usize, _ypos: usize, _cand_mvs: &[MV]) -> (MV, u32) {
520 (ZERO_MV, std::u32::MAX / 2)
524 pub struct MotionEstimator {
528 srch: Box<dyn MVSearch+Send>,
531 impl MotionEstimator {
532 pub fn new() -> Self {
533 let mode = MVSearchMode::default();
541 pub fn get_mode(&self) -> MVSearchMode { self.mode }
542 pub fn set_mode(&mut self, new_mode: MVSearchMode) {
543 if self.mode != new_mode {
544 self.mode = new_mode;
545 self.srch = self.mode.create();
548 pub fn search_mb_p(&mut self, pic: &NAVideoBuffer<u8>, refmb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32) {
549 let mut mv_est = MVEstimator {
550 mv_range: self.range,
551 cutoff_thr: self.thresh,
554 self.srch.search_mb(&mut mv_est, refmb, mb_x, mb_y, cand_mvs)
556 pub fn search_blk8(&mut self, pic: &NAVideoBuffer<u8>, refmb: &RefMBData, xoff: usize, yoff: usize, cand_mvs: &[MV]) -> (MV, u32) {
557 let mut mv_est = MVEstimator {
558 mv_range: self.range,
559 cutoff_thr: self.thresh,
562 self.srch.search_blk8(&mut mv_est, refmb, xoff, yoff, cand_mvs)
566 pub struct SearchB<'a> {
567 ref_p: &'a NAVideoBuffer<u8>,
568 ref_n: &'a NAVideoBuffer<u8>,
577 impl<'a> SearchB<'a> {
578 pub fn new(ref_p: &'a NAVideoBuffer<u8>, ref_n: &'a NAVideoBuffer<u8>, mb_x: usize, mb_y: usize, ratios: [u32; 2]) -> Self {
584 tmp1: RefMBData::new(),
585 tmp2: RefMBData::new(),
586 pred_blk: RefMBData::new(),
589 pub fn search_mb(&mut self, ref_mb: &RefMBData, cand_mvs: [MV; 2]) -> (MV, MV) {
590 let mut best_cand = cand_mvs;
591 let mut best_dist = self.interp_b_dist(ref_mb, best_cand, MAX_DIST);
594 let mut improved = false;
595 for &fmv_add in DIA_PATTERN.iter() {
596 for &bmv_add in DIA_PATTERN.iter() {
597 let cand = [best_cand[0] + fmv_add.from_pixels(),
598 best_cand[1] + bmv_add.from_pixels()];
599 let dist = self.interp_b_dist(ref_mb, cand, best_dist);
600 if dist < best_dist {
612 for &fmv_add in REFINEMENT.iter() {
613 for &bmv_add in REFINEMENT.iter() {
614 let cand = [best_cand[0] + fmv_add, best_cand[1] + bmv_add];
615 let dist = self.interp_b_dist(ref_mb, cand, best_dist);
616 if dist < best_dist {
623 (best_cand[0], best_cand[1])
625 fn interp_b_dist(&mut self, ref_mb: &RefMBData, cand_mv: [MV; 2], cur_best_dist: u32) -> u32 {
626 let [fmv, bmv] = cand_mv;
627 luma_mc(&mut self.tmp1.y, 16, self.ref_p, self.xpos, self.ypos, fmv, true);
628 chroma_mc(&mut self.tmp1.u, 8, self.ref_p, self.xpos / 2, self.ypos / 2, 1, fmv, true);
629 chroma_mc(&mut self.tmp1.v, 8, self.ref_p, self.xpos / 2, self.ypos / 2, 2, fmv, true);
630 luma_mc(&mut self.tmp2.y, 16, self.ref_n, self.xpos, self.ypos, bmv, true);
631 chroma_mc(&mut self.tmp2.u, 8, self.ref_n, self.xpos / 2, self.ypos / 2, 1, bmv, true);
632 chroma_mc(&mut self.tmp2.v, 8, self.ref_n, self.xpos / 2, self.ypos / 2, 2, bmv, true);
633 self.pred_blk.avg(&self.tmp1, self.ratios[0], &self.tmp2, self.ratios[1]);
636 for (dline, sline) in self.pred_blk.y.chunks(16).zip(ref_mb.y.chunks(16)) {
637 dist += sad!(dline, sline);
638 if dist > cur_best_dist {
642 dist += sad!(self.pred_blk.u, ref_mb.u);
643 if dist > cur_best_dist {
646 dist += sad!(self.pred_blk.v, ref_mb.v);
652 macro_rules! hadamard {
653 ($s0:expr, $s1:expr, $s2:expr, $s3:expr, $d0:expr, $d1:expr, $d2:expr, $d3:expr) => {
665 pub struct FrameComplexityEstimate {
666 ref_frm: NAVideoBufferRef<u8>,
667 cur_frm: NAVideoBufferRef<u8>,
668 nxt_frm: NAVideoBufferRef<u8>,
673 impl FrameComplexityEstimate {
674 pub fn new() -> Self {
675 let vinfo = NAVideoInfo::new(24, 24, false, YUV420_FORMAT);
676 let vt = alloc_video_buffer(vinfo, 4).unwrap();
677 let buf = vt.get_vbuf().unwrap();
679 ref_frm: buf.clone(),
680 cur_frm: buf.clone(),
686 pub fn resize(&mut self, width: usize, height: usize) {
687 if width != self.width || height != self.height {
689 self.height = height;
691 let vinfo = NAVideoInfo::new(self.width / 2, self.height / 2, false, YUV420_FORMAT);
692 let vt = alloc_video_buffer(vinfo, 4).unwrap();
693 self.ref_frm = vt.get_vbuf().unwrap();
694 let frm = self.ref_frm.get_data_mut().unwrap();
695 for el in frm.iter_mut() {
698 let vt = alloc_video_buffer(vinfo, 4).unwrap();
699 self.cur_frm = vt.get_vbuf().unwrap();
700 let vt = alloc_video_buffer(vinfo, 4).unwrap();
701 self.nxt_frm = vt.get_vbuf().unwrap();
704 pub fn set_current(&mut self, frm: &NAVideoBuffer<u8>) {
705 Self::downscale(&mut self.cur_frm, frm);
707 pub fn get_complexity(&self, ftype: FrameType) -> u32 {
709 FrameType::I => Self::calculate_i_cplx(&self.cur_frm),
710 FrameType::P => Self::calculate_mv_diff(&self.ref_frm, &self.cur_frm),
714 pub fn decide_b_frame(&mut self, frm1: &NAVideoBuffer<u8>, frm2: &NAVideoBuffer<u8>) -> bool {
715 Self::downscale(&mut self.cur_frm, frm1);
716 Self::downscale(&mut self.nxt_frm, frm2);
717 let diff_ref_cur = Self::calculate_mv_diff(&self.ref_frm, &self.cur_frm);
718 let diff_cur_nxt = Self::calculate_mv_diff(&self.cur_frm, &self.nxt_frm);
720 // simple rule - if complexity ref->cur and cur->next is about the same this should be a B-frame
721 let ddiff = diff_ref_cur.max(diff_cur_nxt) - diff_ref_cur.min(diff_cur_nxt);
726 while (ddiff << order) < diff_ref_cur.min(diff_cur_nxt) {
732 pub fn update_ref(&mut self) {
733 std::mem::swap(&mut self.ref_frm, &mut self.cur_frm);
736 fn add_mv(mb_x: usize, mb_y: usize, mv: MV) -> (usize, usize) {
737 (((mb_x * 16) as isize + (mv.x as isize)) as usize,
738 ((mb_y * 16) as isize + (mv.y as isize)) as usize)
740 fn calculate_i_cplx(frm: &NAVideoBuffer<u8>) -> u32 {
741 let (w, h) = frm.get_dimensions(0);
742 let src = frm.get_data();
743 let stride = frm.get_stride(0);
746 for y in (0..h).step_by(4) {
747 for x in (0..w).step_by(4) {
748 sum += Self::satd_i(src, offset + x, stride, x > 0, y > 0);
750 offset += stride * 4;
754 fn calculate_mv_diff(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>) -> u32 {
755 let (w, h) = ref_frm.get_dimensions(0);
757 for mb_y in 0..(h / 16) {
758 for mb_x in 0..(w / 16) {
759 sum += Self::satd_mb_diff(ref_frm, cur_frm, mb_x, mb_y);
764 fn satd_mb_diff(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>, mb_x: usize, mb_y: usize) -> u32 {
765 let mv = Self::search_mv(ref_frm, cur_frm, mb_x, mb_y);
767 let src0 = ref_frm.get_data();
768 let src1 = cur_frm.get_data();
769 let stride = ref_frm.get_stride(0);
770 let (src_x, src_y) = Self::add_mv(mb_x, mb_y, mv);
771 for y in (0..16).step_by(4) {
772 for x in (0..16).step_by(4) {
773 sum += Self::satd(&src0[src_x + x + (src_y + y) * stride..],
774 &src1[mb_x * 16 + x + (mb_y * 16 + y) * stride..],
780 fn search_mv(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>, mb_x: usize, mb_y: usize) -> MV {
781 let stride = ref_frm.get_stride(0);
782 let (w, h) = ref_frm.get_dimensions(0);
783 let (v_edge, h_edge) = (w - 16, h - 16);
784 let ref_src = ref_frm.get_data();
785 let cur_src = cur_frm.get_data();
786 let cur_src = &cur_src[mb_x * 16 + mb_y * 16 * stride..];
788 let mut best_mv = ZERO_MV;
789 let mut best_dist = Self::sad(cur_src, ref_src, mb_x, mb_y, stride, best_mv);
794 for step in (0..=2).rev() {
795 let mut changed = true;
798 for &mv in DIA_PATTERN[1..].iter() {
799 let cand_mv = best_mv + mv.scale(1 << step);
800 let (cx, cy) = Self::add_mv(mb_x, mb_y, cand_mv);
801 if cx > v_edge || cy > h_edge {
804 let cand_dist = Self::sad(cur_src, ref_src, mb_x, mb_y, stride, cand_mv);
805 if cand_dist < best_dist {
806 best_dist = cand_dist;
818 fn sad(cur_src: &[u8], src: &[u8], mb_x: usize, mb_y: usize, stride: usize, mv: MV) -> u32 {
819 let (src_x, src_y) = Self::add_mv(mb_x, mb_y, mv);
821 for (line1, line2) in cur_src.chunks(stride).zip(src[src_x + src_y * stride..].chunks(stride)).take(16) {
822 sum += line1[..16].iter().zip(line2[..16].iter()).fold(0u32,
823 |acc, (&a, &b)| acc + u32::from(a.max(b) - a.min(b)) * u32::from(a.max(b) - a.min(b)));
827 fn satd_i(src: &[u8], mut offset: usize, stride: usize, has_left: bool, has_top: bool) -> u32 {
828 let mut diffs = [0; 16];
829 match (has_left, has_top) {
831 for row in diffs.chunks_exact_mut(4) {
832 let mut left = i16::from(src[offset - 1]);
833 let mut tl = i16::from(src[offset - stride - 1]);
834 for (x, dst) in row.iter_mut().enumerate() {
835 let cur = i16::from(src[offset + x]);
836 let top = i16::from(src[offset + x - stride]);
838 *dst = cur - (top + left + tl - top.min(left).min(tl) - top.max(left).max(tl));
848 for (dst, (left, cur)) in diffs.chunks_exact_mut(4).zip(
849 src[offset - 1..].chunks(stride).zip(src[offset..].chunks(stride))) {
850 for (dst, (&left, &cur)) in dst.iter_mut().zip(left.iter().zip(cur.iter())) {
851 *dst = i16::from(cur) - i16::from(left);
856 for (dst, (top, cur)) in diffs.chunks_exact_mut(4).zip(
857 src[offset - stride..].chunks(stride).zip(src[offset..].chunks(stride))) {
858 for (dst, (&top, &cur)) in dst.iter_mut().zip(top.iter().zip(cur.iter())) {
859 *dst = i16::from(cur) - i16::from(top);
864 for (dst, src) in diffs.chunks_exact_mut(4).zip(src[offset..].chunks(stride)) {
865 for (dst, &src) in dst.iter_mut().zip(src.iter()) {
866 *dst = i16::from(src) - 128;
871 for row in diffs.chunks_exact_mut(4) {
872 hadamard!(row[0], row[1], row[2], row[3], row[0], row[1], row[2], row[3]);
875 hadamard!(diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12],
876 diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12]);
878 diffs.iter().fold(0u32, |acc, x| acc + (x.unsigned_abs() as u32))
880 fn satd(src0: &[u8], src1: &[u8], stride: usize) -> u32 {
881 let mut diffs = [0; 16];
882 for (dst, (src0, src1)) in diffs.chunks_exact_mut(4).zip(
883 src0.chunks(stride).zip(src1.chunks(stride))) {
884 hadamard!(i16::from(src0[0]) - i16::from(src1[0]),
885 i16::from(src0[1]) - i16::from(src1[1]),
886 i16::from(src0[2]) - i16::from(src1[2]),
887 i16::from(src0[3]) - i16::from(src1[3]),
888 dst[0], dst[1], dst[2], dst[3]);
891 hadamard!(diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12],
892 diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12]);
894 diffs.iter().fold(0u32, |acc, x| acc + (x.unsigned_abs() as u32))
896 fn downscale(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>) {
897 let dst = NASimpleVideoFrame::from_video_buf(dst).unwrap();
898 let sdata = src.get_data();
900 let cur_w = dst.width[plane];
901 let cur_h = dst.height[plane];
902 let doff = dst.offset[plane];
903 let soff = src.get_offset(plane);
904 let dstride = dst.stride[plane];
905 let sstride = src.get_stride(plane);
906 for (dline, sstrip) in dst.data[doff..].chunks_exact_mut(dstride).zip(
907 sdata[soff..].chunks_exact(sstride * 2)).take(cur_h) {
908 let (line0, line1) = sstrip.split_at(sstride);
909 for (dst, (src0, src1)) in dline.iter_mut().zip(
910 line0.chunks_exact(2).zip(line1.chunks_exact(2))).take(cur_w) {
911 *dst = ((u16::from(src0[0]) + u16::from(src0[1]) +
912 u16::from(src1[0]) + u16::from(src1[1]) + 2) >> 2) as u8;