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)]
7 pub enum MVSearchMode {
15 pub const fn get_possible_modes() -> &'static [&'static str] {
16 &["diamond", "hexagon", "umh"]
18 fn create(self) -> Box<dyn MVSearch+Send> {
20 MVSearchMode::Dummy => Box::new(DummySearcher{}),
21 MVSearchMode::Diamond => Box::new(DiaSearch::new()),
22 MVSearchMode::Hexagon => Box::new(HexSearch::new()),
23 MVSearchMode::UMH => Box::new(UnevenHexSearch::new()),
28 impl Default for MVSearchMode {
29 fn default() -> Self { MVSearchMode::Hexagon }
32 impl std::fmt::Display for MVSearchMode {
33 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
35 MVSearchMode::Diamond => write!(f, "diamond"),
36 MVSearchMode::Hexagon => write!(f, "hexagon"),
37 MVSearchMode::UMH => write!(f, "umh"),
38 MVSearchMode::Dummy => write!(f, "dummy"),
43 impl FromStr for MVSearchMode {
45 fn from_str(s: &str) -> Result<Self, Self::Err> {
47 "diamond" => Ok(MVSearchMode::Diamond),
48 "hexagon" => Ok(MVSearchMode::Hexagon),
49 "umh" => Ok(MVSearchMode::UMH),
50 "dummy" => Ok(MVSearchMode::Dummy),
56 const MAX_DIST: u32 = std::u32::MAX;
57 const DIST_THRESH: u32 = 256;
60 fn from_pixels(self) -> Self;
63 impl FromPixels for MV {
64 fn from_pixels(self) -> MV {
65 MV { x: self.x * 4, y: self.y * 4 }
69 const DIA_PATTERN: [MV; 9] = [
81 const HEX_PATTERN: [MV; 7] = [
91 const REFINEMENT: [MV; 4] = [
98 macro_rules! search_template {
99 ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident, $threshold: expr) => ({
100 search_template!($self, $mv_est, $cur_blk, $mb_x, $mb_y, $sad_func, $threshold, ZERO_MV, MAX_DIST, true)
102 ($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) => ({
103 let mut best_dist = $best_dist;
104 let mut best_mv = $start_mv;
112 let mut cur_best_dist = best_dist;
113 for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
114 if *dist == MAX_DIST {
115 *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point.from_pixels(), cur_best_dist);
116 cur_best_dist = cur_best_dist.min(*dist);
117 if *dist <= $threshold {
122 min_dist = $self.dist[0];
124 for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
128 if dist <= $threshold {
133 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 {
136 best_dist = min_dist;
137 $self.update($self.steps[min_idx]);
139 best_dist = min_dist;
140 best_mv = $self.point[min_idx];
141 if best_dist <= $threshold {
142 return (best_mv.from_pixels(), best_dist);
144 for &step in REFINEMENT.iter() {
145 let mv = best_mv + step;
146 let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv.from_pixels(), MAX_DIST);
147 if best_dist > dist {
152 best_mv = best_mv.from_pixels();
153 if best_dist <= $threshold {
154 return (best_mv, best_dist);
159 $self.set_new_point(best_mv, best_dist);
161 let mut cur_best_dist = best_dist;
162 for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) {
163 if *dist == MAX_DIST {
164 *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point, cur_best_dist);
165 cur_best_dist = cur_best_dist.min(*dist);
166 if *dist <= $threshold {
171 min_dist = $self.dist[0];
173 for (i, &dist) in $self.dist.iter().enumerate().skip(1) {
177 if dist <= $threshold {
182 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 {
185 best_dist = min_dist;
186 $self.update($self.steps[min_idx]);
188 best_dist = min_dist;
189 best_mv = $self.point[min_idx];
190 if best_dist <= $threshold {
191 return (best_mv, best_dist);
193 for &step in REFINEMENT.iter() {
194 let mv = best_mv + step;
195 let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv, MAX_DIST);
196 if best_dist > dist {
205 macro_rules! pattern_search {
206 ($struct_name: ident, $patterns: expr) => {
207 pub struct $struct_name {
208 point: [MV; $patterns.len()],
209 dist: [u32; $patterns.len()],
210 steps: &'static [MV; $patterns.len()],
214 pub fn new() -> Self {
217 dist: [MAX_DIST; $patterns.len()],
221 fn reset(&mut self) {
222 self.point = $patterns;
223 self.dist = [MAX_DIST; $patterns.len()];
225 fn set_new_point(&mut self, start: MV, dist: u32) {
226 for (dst, &src) in self.point.iter_mut().zip(self.steps.iter()) {
229 self.dist = [MAX_DIST; $patterns.len()];
232 fn update(&mut self, step: MV) {
233 let mut new_point = self.point;
234 let mut new_dist = [MAX_DIST; $patterns.len()];
236 for point in new_point.iter_mut() {
240 for (new_point, new_dist) in new_point.iter_mut().zip(new_dist.iter_mut()) {
241 for (&old_point, &old_dist) in self.point.iter().zip(self.dist.iter()) {
242 if *new_point == old_point {
243 *new_dist = old_dist;
248 self.point = new_point;
249 self.dist = new_dist;
253 impl MVSearch for $struct_name {
254 fn search_mb(&mut self, mv_est: &mut MVEstimator, cur_mb: &RefMBData, mb_x: usize, mb_y: usize, _cand_mvs: &[MV]) -> (MV, u32) {
255 search_template!(self, mv_est, cur_mb, mb_x, mb_y, sad_mb, DIST_THRESH)
257 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, _cand_mvs: &[MV]) -> (MV, u32) {
258 search_template!(self, mv_est, ref_blk, xpos, ypos, sad_blk8, DIST_THRESH / 4)
264 pattern_search!(DiaSearch, DIA_PATTERN);
265 pattern_search!(HexSearch, HEX_PATTERN);
267 const LARGE_HEX_PATTERN: [MV; 16] = [
286 const UNSYMM_CROSS: [MV; 4] = [
294 struct UniqueSet<T:Copy+Default> {
299 impl<T:Copy+Default+PartialEq> UniqueSet<T> {
300 fn new() -> Self { Self::default() }
301 fn clear(&mut self) { self.count = 0; }
302 fn get_list(&self) -> &[T] { &self.list[..self.count] }
303 fn add(&mut self, val: T) {
304 if self.count < self.list.len() && !self.get_list().contains(&val) {
305 self.list[self.count] = val;
312 fn scale(self, scale: i16) -> Self;
313 fn is_in_range(self, range: i16) -> bool;
317 fn scale(self, scale: i16) -> MV {
318 MV { x: self.x * scale, y: self.y * scale }
320 fn is_in_range(self, range: i16) -> bool {
321 self.x.abs() <= range && self.y.abs() <= range
325 macro_rules! single_search_step {
326 ($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) => {{
327 let mut best_mv = $start;
328 let mut best_dist = $best_dist;
329 for point in $pattern.iter() {
330 let mv = point.scale($scale) + $start;
331 if !mv.is_in_range($mv_est.mv_range * 4) {
334 let dist = $mv_est.$sad_func($ref_blk, $xpos, $ypos, mv, best_dist);
335 if dist < best_dist {
338 if best_dist < $dist_thr {
343 (best_mv, best_dist, best_mv != $start)
347 struct UnevenHexSearch {
348 mv_list: UniqueSet<MV>,
351 impl UnevenHexSearch {
354 mv_list: UniqueSet::new(),
357 fn get_cand_mv(&mut self, cand_mvs: &[MV]) -> MV {
358 self.mv_list.clear();
359 for &mv in cand_mvs.iter() {
360 self.mv_list.add(mv);
362 match self.mv_list.count {
363 1 => self.mv_list.list[0],
364 3 => MV::pred(self.mv_list.list[0], self.mv_list.list[1], self.mv_list.list[2]),
366 let sum = self.mv_list.get_list().iter().fold((0i32, 0i32),
367 |acc, mv| (acc.0 + i32::from(mv.x), acc.1 + i32::from(mv.y)));
368 MV {x: (sum.0 / (self.mv_list.count as i32)) as i16,
369 y: (sum.1 / (self.mv_list.count as i32)) as i16}
375 macro_rules! umh_search_template {
376 ($cand_mv:expr, $cutoff:expr, $mv_est:expr, $sad_func:ident, $ref_blk:expr, $xpos:expr, $ypos:expr) => {{
377 let cand_mv = $cand_mv;
378 let best_dist = $mv_est.$sad_func($ref_blk, $xpos, $ypos, cand_mv, MAX_DIST);
379 if best_dist < $cutoff {
380 return (cand_mv, best_dist);
383 // step 1 - small refinement search
384 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);
385 if best_dist < $cutoff {
386 return (cand_mv, best_dist);
389 // step 2 - unsymmetrical cross search
391 let (mv, dist, changed) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, UNSYMM_CROSS, 4, $cutoff);
397 if best_dist < $cutoff {
402 // step 3 - multi-hexagon grid search
405 let (mv, dist, changed) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, LARGE_HEX_PATTERN, scale, $cutoff);
411 if best_dist < $cutoff {
416 // step 4 - final hexagon search
417 let (cand_mv, best_dist, _) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, HEX_PATTERN, 1, $cutoff);
418 if best_dist > $cutoff {
419 let (mv, dist, _) = single_search_step!(cand_mv, best_dist, $mv_est, $sad_func, $ref_blk, $xpos, $ypos, DIA_PATTERN, 1, $cutoff);
427 impl MVSearch for UnevenHexSearch {
428 fn search_mb(&mut self, mv_est: &mut MVEstimator, cur_mb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32) {
429 let cand_mv = self.get_cand_mv(cand_mvs);
430 let cutoff = mv_est.cutoff_thr;
431 umh_search_template!(cand_mv, cutoff, mv_est, sad_mb, cur_mb, mb_x, mb_y)
433 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, cand_mvs: &[MV]) -> (MV, u32) {
434 let cand_mv = self.get_cand_mv(cand_mvs);
435 let cutoff = mv_est.cutoff_thr / 4;
436 umh_search_template!(cand_mv, cutoff, mv_est, sad_blk8, ref_blk, xpos, ypos)
440 struct MVEstimator<'a> {
441 pic: &'a NAVideoBuffer<u8>,
447 ($src1:expr, $src2:expr) => {
448 $src1.iter().zip($src2.iter()).fold(0u32, |acc, (&a, &b)|
449 acc + (((i32::from(a) - i32::from(b)) * (i32::from(a) - i32::from(b))) as u32))
453 impl<'a> MVEstimator<'a> {
454 fn sad_mb(&self, ref_mb: &RefMBData, mb_x: usize, mb_y: usize, mv: MV, cur_best_dist: u32) -> u32 {
455 let mut dst = RefMBData::new();
456 luma_mc(&mut dst.y, 16, self.pic, mb_x * 16, mb_y * 16, mv, true);
459 for (dline, sline) in dst.y.chunks(16).zip(ref_mb.y.chunks(16)) {
460 dist += sad!(dline, sline);
461 if dist > cur_best_dist {
465 chroma_mc(&mut dst.u, 8, self.pic, mb_x * 8, mb_y * 8, 1, mv, true);
466 dist += sad!(dst.u, ref_mb.u);
467 if dist > cur_best_dist {
470 chroma_mc(&mut dst.v, 8, self.pic, mb_x * 8, mb_y * 8, 2, mv, true);
471 dist += sad!(dst.v, ref_mb.v);
475 fn sad_blk8(&self, ref_mb: &RefMBData, xpos: usize, ypos: usize, mv: MV, cur_best_dist: u32) -> u32 {
476 let mut cur_y = [0; 64];
477 let mut cur_u = [0; 16];
478 let mut cur_v = [0; 16];
482 let y_off = (xpos & 8) + (ypos & 8) * 16;
483 luma_mc(&mut cur_y, 8, self.pic, xpos, ypos, mv, false);
484 for (dline, sline) in cur_y.chunks(8).zip(ref_mb.y[y_off..].chunks(16)) {
485 dist += sad!(dline, sline);
486 if dist > cur_best_dist {
491 let c_off = (xpos & 8) / 2 + (ypos & 8) * 4;
492 chroma_mc(&mut cur_u, 4, self.pic, xpos / 2, ypos / 2, 1, mv, false);
493 for (dline, sline) in cur_u.chunks(4).zip(ref_mb.u[c_off..].chunks(8)) {
494 dist += sad!(dline, sline);
495 if dist > cur_best_dist {
499 chroma_mc(&mut cur_v, 4, self.pic, xpos / 2, ypos / 2, 2, mv, false);
500 for (dline, sline) in cur_v.chunks(4).zip(ref_mb.v[c_off..].chunks(8)) {
501 dist += sad!(dline, sline);
502 if dist > cur_best_dist {
512 fn search_mb(&mut self, mv_est: &mut MVEstimator, ref_mb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32);
513 fn search_blk8(&mut self, mv_est: &mut MVEstimator, ref_blk: &RefMBData, xpos: usize, ypos: usize, cand_mvs: &[MV]) -> (MV, u32);
516 struct DummySearcher {}
518 impl MVSearch for DummySearcher {
519 fn search_mb(&mut self, _mv_est: &mut MVEstimator, _ref_mb: &RefMBData, _mb_x: usize, _mb_y: usize, _cand_mvs: &[MV]) -> (MV, u32) {
520 (ZERO_MV, std::u32::MAX / 2)
522 fn search_blk8(&mut self, _mv_est: &mut MVEstimator, _ref_mb: &RefMBData, _xpos: usize, _ypos: usize, _cand_mvs: &[MV]) -> (MV, u32) {
523 (ZERO_MV, std::u32::MAX / 2)
527 pub struct MotionEstimator {
531 srch: Box<dyn MVSearch+Send>,
534 impl MotionEstimator {
535 pub fn new() -> Self {
536 let mode = MVSearchMode::default();
544 pub fn get_mode(&self) -> MVSearchMode { self.mode }
545 pub fn set_mode(&mut self, new_mode: MVSearchMode) {
546 if self.mode != new_mode {
547 self.mode = new_mode;
548 self.srch = self.mode.create();
551 pub fn search_mb_p(&mut self, pic: &NAVideoBuffer<u8>, refmb: &RefMBData, mb_x: usize, mb_y: usize, cand_mvs: &[MV]) -> (MV, u32) {
552 let mut mv_est = MVEstimator {
553 mv_range: self.range,
554 cutoff_thr: self.thresh,
557 self.srch.search_mb(&mut mv_est, refmb, mb_x, mb_y, cand_mvs)
559 pub fn search_blk8(&mut self, pic: &NAVideoBuffer<u8>, refmb: &RefMBData, xoff: usize, yoff: usize, cand_mvs: &[MV]) -> (MV, u32) {
560 let mut mv_est = MVEstimator {
561 mv_range: self.range,
562 cutoff_thr: self.thresh,
565 self.srch.search_blk8(&mut mv_est, refmb, xoff, yoff, cand_mvs)
569 pub struct SearchB<'a> {
570 ref_p: &'a NAVideoBuffer<u8>,
571 ref_n: &'a NAVideoBuffer<u8>,
580 impl<'a> SearchB<'a> {
581 pub fn new(ref_p: &'a NAVideoBuffer<u8>, ref_n: &'a NAVideoBuffer<u8>, mb_x: usize, mb_y: usize, ratios: [u32; 2]) -> Self {
587 tmp1: RefMBData::new(),
588 tmp2: RefMBData::new(),
589 pred_blk: RefMBData::new(),
592 pub fn search_mb(&mut self, ref_mb: &RefMBData, cand_mvs: [MV; 2]) -> (MV, MV) {
593 let mut best_cand = cand_mvs;
594 let mut best_dist = self.interp_b_dist(ref_mb, best_cand, MAX_DIST);
597 let mut improved = false;
598 for &fmv_add in DIA_PATTERN.iter() {
599 for &bmv_add in DIA_PATTERN.iter() {
600 let cand = [best_cand[0] + fmv_add.from_pixels(),
601 best_cand[1] + bmv_add.from_pixels()];
602 let dist = self.interp_b_dist(ref_mb, cand, best_dist);
603 if dist < best_dist {
615 for &fmv_add in REFINEMENT.iter() {
616 for &bmv_add in REFINEMENT.iter() {
617 let cand = [best_cand[0] + fmv_add, best_cand[1] + bmv_add];
618 let dist = self.interp_b_dist(ref_mb, cand, best_dist);
619 if dist < best_dist {
626 (best_cand[0], best_cand[1])
628 fn interp_b_dist(&mut self, ref_mb: &RefMBData, cand_mv: [MV; 2], cur_best_dist: u32) -> u32 {
629 let [fmv, bmv] = cand_mv;
630 luma_mc(&mut self.tmp1.y, 16, self.ref_p, self.xpos, self.ypos, fmv, true);
631 chroma_mc(&mut self.tmp1.u, 8, self.ref_p, self.xpos / 2, self.ypos / 2, 1, fmv, true);
632 chroma_mc(&mut self.tmp1.v, 8, self.ref_p, self.xpos / 2, self.ypos / 2, 2, fmv, true);
633 luma_mc(&mut self.tmp2.y, 16, self.ref_n, self.xpos, self.ypos, bmv, true);
634 chroma_mc(&mut self.tmp2.u, 8, self.ref_n, self.xpos / 2, self.ypos / 2, 1, bmv, true);
635 chroma_mc(&mut self.tmp2.v, 8, self.ref_n, self.xpos / 2, self.ypos / 2, 2, bmv, true);
636 self.pred_blk.avg(&self.tmp1, self.ratios[0], &self.tmp2, self.ratios[1]);
639 for (dline, sline) in self.pred_blk.y.chunks(16).zip(ref_mb.y.chunks(16)) {
640 dist += sad!(dline, sline);
641 if dist > cur_best_dist {
645 dist += sad!(self.pred_blk.u, ref_mb.u);
646 if dist > cur_best_dist {
649 dist += sad!(self.pred_blk.v, ref_mb.v);
655 macro_rules! hadamard {
656 ($s0:expr, $s1:expr, $s2:expr, $s3:expr, $d0:expr, $d1:expr, $d2:expr, $d3:expr) => {
668 pub struct FrameComplexityEstimate {
669 ref_frm: NAVideoBufferRef<u8>,
670 cur_frm: NAVideoBufferRef<u8>,
671 nxt_frm: NAVideoBufferRef<u8>,
676 impl FrameComplexityEstimate {
677 pub fn new() -> Self {
678 let vinfo = NAVideoInfo::new(24, 24, false, YUV420_FORMAT);
679 let vt = alloc_video_buffer(vinfo, 4).unwrap();
680 let buf = vt.get_vbuf().unwrap();
682 ref_frm: buf.clone(),
683 cur_frm: buf.clone(),
689 pub fn resize(&mut self, width: usize, height: usize) {
690 if width != self.width || height != self.height {
692 self.height = height;
694 let vinfo = NAVideoInfo::new(self.width / 2, self.height / 2, false, YUV420_FORMAT);
695 let vt = alloc_video_buffer(vinfo, 4).unwrap();
696 self.ref_frm = vt.get_vbuf().unwrap();
697 let frm = self.ref_frm.get_data_mut().unwrap();
698 for el in frm.iter_mut() {
701 let vt = alloc_video_buffer(vinfo, 4).unwrap();
702 self.cur_frm = vt.get_vbuf().unwrap();
703 let vt = alloc_video_buffer(vinfo, 4).unwrap();
704 self.nxt_frm = vt.get_vbuf().unwrap();
707 pub fn set_current(&mut self, frm: &NAVideoBuffer<u8>) {
708 Self::downscale(&mut self.cur_frm, frm);
710 pub fn get_complexity(&self, ftype: FrameType) -> u32 {
712 FrameType::I => Self::calculate_i_cplx(&self.cur_frm),
713 FrameType::P => Self::calculate_mv_diff(&self.ref_frm, &self.cur_frm),
717 pub fn decide_b_frame(&mut self, frm1: &NAVideoBuffer<u8>, frm2: &NAVideoBuffer<u8>) -> bool {
718 Self::downscale(&mut self.cur_frm, frm1);
719 Self::downscale(&mut self.nxt_frm, frm2);
720 let diff_ref_cur = Self::calculate_mv_diff(&self.ref_frm, &self.cur_frm);
721 let diff_cur_nxt = Self::calculate_mv_diff(&self.cur_frm, &self.nxt_frm);
723 // simple rule - if complexity ref->cur and cur->next is about the same this should be a B-frame
724 let ddiff = diff_ref_cur.max(diff_cur_nxt) - diff_ref_cur.min(diff_cur_nxt);
729 while (ddiff << order) < diff_ref_cur.min(diff_cur_nxt) {
735 pub fn update_ref(&mut self) {
736 std::mem::swap(&mut self.ref_frm, &mut self.cur_frm);
739 fn add_mv(mb_x: usize, mb_y: usize, mv: MV) -> (usize, usize) {
740 (((mb_x * 16) as isize + (mv.x as isize)) as usize,
741 ((mb_y * 16) as isize + (mv.y as isize)) as usize)
743 fn calculate_i_cplx(frm: &NAVideoBuffer<u8>) -> u32 {
744 let (w, h) = frm.get_dimensions(0);
745 let src = frm.get_data();
746 let stride = frm.get_stride(0);
749 for y in (0..h).step_by(4) {
750 for x in (0..w).step_by(4) {
751 sum += Self::satd_i(src, offset + x, stride, x > 0, y > 0);
753 offset += stride * 4;
757 fn calculate_mv_diff(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>) -> u32 {
758 let (w, h) = ref_frm.get_dimensions(0);
760 for mb_y in 0..(h / 16) {
761 for mb_x in 0..(w / 16) {
762 sum += Self::satd_mb_diff(ref_frm, cur_frm, mb_x, mb_y);
767 fn satd_mb_diff(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>, mb_x: usize, mb_y: usize) -> u32 {
768 let mv = Self::search_mv(ref_frm, cur_frm, mb_x, mb_y);
770 let src0 = ref_frm.get_data();
771 let src1 = cur_frm.get_data();
772 let stride = ref_frm.get_stride(0);
773 let (src_x, src_y) = Self::add_mv(mb_x, mb_y, mv);
774 for y in (0..16).step_by(4) {
775 for x in (0..16).step_by(4) {
776 sum += Self::satd(&src0[src_x + x + (src_y + y) * stride..],
777 &src1[mb_x * 16 + x + (mb_y * 16 + y) * stride..],
783 fn search_mv(ref_frm: &NAVideoBuffer<u8>, cur_frm: &NAVideoBuffer<u8>, mb_x: usize, mb_y: usize) -> MV {
784 let stride = ref_frm.get_stride(0);
785 let (w, h) = ref_frm.get_dimensions(0);
786 let (v_edge, h_edge) = (w - 16, h - 16);
787 let ref_src = ref_frm.get_data();
788 let cur_src = cur_frm.get_data();
789 let cur_src = &cur_src[mb_x * 16 + mb_y * 16 * stride..];
791 let mut best_mv = ZERO_MV;
792 let mut best_dist = Self::sad(cur_src, ref_src, mb_x, mb_y, stride, best_mv);
797 for step in (0..=2).rev() {
798 let mut changed = true;
801 for &mv in DIA_PATTERN[1..].iter() {
802 let cand_mv = best_mv + mv.scale(1 << step);
803 let (cx, cy) = Self::add_mv(mb_x, mb_y, cand_mv);
804 if cx > v_edge || cy > h_edge {
807 let cand_dist = Self::sad(cur_src, ref_src, mb_x, mb_y, stride, cand_mv);
808 if cand_dist < best_dist {
809 best_dist = cand_dist;
821 fn sad(cur_src: &[u8], src: &[u8], mb_x: usize, mb_y: usize, stride: usize, mv: MV) -> u32 {
822 let (src_x, src_y) = Self::add_mv(mb_x, mb_y, mv);
824 for (line1, line2) in cur_src.chunks(stride).zip(src[src_x + src_y * stride..].chunks(stride)).take(16) {
825 sum += line1[..16].iter().zip(line2[..16].iter()).fold(0u32,
826 |acc, (&a, &b)| acc + u32::from(a.max(b) - a.min(b)) * u32::from(a.max(b) - a.min(b)));
830 fn satd_i(src: &[u8], mut offset: usize, stride: usize, has_left: bool, has_top: bool) -> u32 {
831 let mut diffs = [0; 16];
832 match (has_left, has_top) {
834 for row in diffs.chunks_exact_mut(4) {
835 let mut left = i16::from(src[offset - 1]);
836 let mut tl = i16::from(src[offset - stride - 1]);
837 for (x, dst) in row.iter_mut().enumerate() {
838 let cur = i16::from(src[offset + x]);
839 let top = i16::from(src[offset + x - stride]);
841 *dst = cur - (top + left + tl - top.min(left).min(tl) - top.max(left).max(tl));
851 for (dst, (left, cur)) in diffs.chunks_exact_mut(4).zip(
852 src[offset - 1..].chunks(stride).zip(src[offset..].chunks(stride))) {
853 for (dst, (&left, &cur)) in dst.iter_mut().zip(left.iter().zip(cur.iter())) {
854 *dst = i16::from(cur) - i16::from(left);
859 for (dst, (top, cur)) in diffs.chunks_exact_mut(4).zip(
860 src[offset - stride..].chunks(stride).zip(src[offset..].chunks(stride))) {
861 for (dst, (&top, &cur)) in dst.iter_mut().zip(top.iter().zip(cur.iter())) {
862 *dst = i16::from(cur) - i16::from(top);
867 for (dst, src) in diffs.chunks_exact_mut(4).zip(src[offset..].chunks(stride)) {
868 for (dst, &src) in dst.iter_mut().zip(src.iter()) {
869 *dst = i16::from(src) - 128;
874 for row in diffs.chunks_exact_mut(4) {
875 hadamard!(row[0], row[1], row[2], row[3], row[0], row[1], row[2], row[3]);
878 hadamard!(diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12],
879 diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12]);
881 diffs.iter().fold(0u32, |acc, x| acc + (x.abs() as u32))
883 fn satd(src0: &[u8], src1: &[u8], stride: usize) -> u32 {
884 let mut diffs = [0; 16];
885 for (dst, (src0, src1)) in diffs.chunks_exact_mut(4).zip(
886 src0.chunks(stride).zip(src1.chunks(stride))) {
887 hadamard!(i16::from(src0[0]) - i16::from(src1[0]),
888 i16::from(src0[1]) - i16::from(src1[1]),
889 i16::from(src0[2]) - i16::from(src1[2]),
890 i16::from(src0[3]) - i16::from(src1[3]),
891 dst[0], dst[1], dst[2], dst[3]);
894 hadamard!(diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12],
895 diffs[i], diffs[i + 4], diffs[i + 8], diffs[i + 12]);
897 diffs.iter().fold(0u32, |acc, x| acc + (x.abs() as u32))
899 fn downscale(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>) {
900 let dst = NASimpleVideoFrame::from_video_buf(dst).unwrap();
901 let sdata = src.get_data();
903 let cur_w = dst.width[plane];
904 let cur_h = dst.height[plane];
905 let doff = dst.offset[plane];
906 let soff = src.get_offset(plane);
907 let dstride = dst.stride[plane];
908 let sstride = src.get_stride(plane);
909 for (dline, sstrip) in dst.data[doff..].chunks_exact_mut(dstride).zip(
910 sdata[soff..].chunks_exact(sstride * 2)).take(cur_h) {
911 let (line0, line1) = sstrip.split_at(sstride);
912 for (dst, (src0, src1)) in dline.iter_mut().zip(
913 line0.chunks_exact(2).zip(line1.chunks_exact(2))).take(cur_w) {
914 *dst = ((u16::from(src0[0]) + u16::from(src0[1]) +
915 u16::from(src1[0]) + u16::from(src1[1]) + 2) >> 2) as u8;