| 1 | use super::*; |
| 2 | use super::kernel::Kernel; |
| 3 | |
| 4 | struct NNResampler {} |
| 5 | |
| 6 | impl NNResampler { |
| 7 | fn new() -> Self { Self{} } |
| 8 | } |
| 9 | |
| 10 | #[allow(clippy::comparison_chain)] |
| 11 | fn scale_line<T:Copy>(src: &[T], dst: &mut [T], src_w: usize, dst_w: usize) { |
| 12 | if src_w == dst_w { |
| 13 | (&mut dst[..dst_w]).copy_from_slice(&src[..dst_w]); |
| 14 | } else if src_w < dst_w { |
| 15 | if dst_w % src_w == 0 { |
| 16 | let step = dst_w / src_w; |
| 17 | for (out, srcv) in dst.chunks_exact_mut(step).take(src_w).zip(src.iter()) { |
| 18 | for el in out.iter_mut() { |
| 19 | *el = *srcv; |
| 20 | } |
| 21 | } |
| 22 | } else { |
| 23 | let mut pos = 0; |
| 24 | for out in dst.iter_mut().take(dst_w) { |
| 25 | *out = src[pos / dst_w]; |
| 26 | pos += src_w; |
| 27 | } |
| 28 | } |
| 29 | } else { |
| 30 | if dst_w % src_w == 0 { |
| 31 | let step = src_w / dst_w; |
| 32 | for (out, srcv) in dst.iter_mut().take(dst_w).zip(src.iter().step_by(step)) { |
| 33 | *out = *srcv; |
| 34 | } |
| 35 | } else { |
| 36 | let mut pos = 0; |
| 37 | for out in dst.iter_mut().take(dst_w) { |
| 38 | *out = src[pos / dst_w]; |
| 39 | pos += src_w; |
| 40 | } |
| 41 | } |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | fn fill_plane<T: Copy>(dst: &mut [T], w: usize, h: usize, stride: usize, val: T) { |
| 46 | for row in dst.chunks_mut(stride).take(h) { |
| 47 | for el in row.iter_mut().take(w) { |
| 48 | *el = val; |
| 49 | } |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | macro_rules! scale_loop { |
| 54 | ($sbuf:expr, $dbuf:expr) => { |
| 55 | let fmt = $sbuf.get_info().get_format(); |
| 56 | let ncomp = fmt.get_num_comp(); |
| 57 | for comp in 0..ncomp { |
| 58 | let istride = $sbuf.get_stride(comp); |
| 59 | let dstride = $dbuf.get_stride(comp); |
| 60 | let (sw, sh) = $sbuf.get_dimensions(comp); |
| 61 | let (dw, dh) = $dbuf.get_dimensions(comp); |
| 62 | let ioff = $sbuf.get_offset(comp); |
| 63 | let mut doff = $dbuf.get_offset(comp); |
| 64 | let src = $sbuf.get_data(); |
| 65 | let dst = $dbuf.get_data_mut().unwrap(); |
| 66 | for y in 0..dh { |
| 67 | let sy = y * sh / dh; |
| 68 | let soff = ioff + sy * istride; |
| 69 | scale_line(&src[soff..], &mut dst[doff..], sw, dw); |
| 70 | doff += dstride; |
| 71 | } |
| 72 | } |
| 73 | let dfmt = $dbuf.get_info().get_format(); |
| 74 | let ndcomp = dfmt.get_num_comp(); |
| 75 | if ndcomp > ncomp { |
| 76 | if !fmt.alpha && dfmt.alpha { |
| 77 | let acomp = ndcomp - 1; |
| 78 | let dstride = $dbuf.get_stride(acomp); |
| 79 | let (dw, dh) = $dbuf.get_dimensions(acomp); |
| 80 | let doff = $dbuf.get_offset(acomp); |
| 81 | let dst = $dbuf.get_data_mut().unwrap(); |
| 82 | fill_plane(&mut dst[doff..], dw, dh, dstride, 0); |
| 83 | } |
| 84 | if fmt.model.is_yuv() && ((!fmt.alpha && ncomp == 1) || (fmt.alpha && ncomp == 2)) && ndcomp >= 3 { |
| 85 | let uval = 1 << (dfmt.comp_info[1].unwrap().depth - 1); |
| 86 | let vval = 1 << (dfmt.comp_info[2].unwrap().depth - 1); |
| 87 | |
| 88 | let ustride = $dbuf.get_stride(1); |
| 89 | let vstride = $dbuf.get_stride(2); |
| 90 | let (uw, uh) = $dbuf.get_dimensions(1); |
| 91 | let (vw, vh) = $dbuf.get_dimensions(2); |
| 92 | let uoff = $dbuf.get_offset(1); |
| 93 | let voff = $dbuf.get_offset(2); |
| 94 | let dst = $dbuf.get_data_mut().unwrap(); |
| 95 | fill_plane(&mut dst[uoff..], uw, uh, ustride, uval); |
| 96 | fill_plane(&mut dst[voff..], vw, vh, vstride, vval); |
| 97 | } |
| 98 | } |
| 99 | }; |
| 100 | } |
| 101 | |
| 102 | impl Kernel for NNResampler { |
| 103 | fn init(&mut self, in_fmt: &ScaleInfo, dest_fmt: &ScaleInfo, _options: &[(String, String)]) -> ScaleResult<NABufferType> { |
| 104 | let res = alloc_video_buffer(NAVideoInfo::new(dest_fmt.width, dest_fmt.height, false, in_fmt.fmt), 3); |
| 105 | if res.is_err() { return Err(ScaleError::AllocError); } |
| 106 | Ok(res.unwrap()) |
| 107 | } |
| 108 | fn process(&mut self, pic_in: &NABufferType, pic_out: &mut NABufferType) { |
| 109 | if let (Some(ref sbuf), Some(ref mut dbuf)) = (pic_in.get_vbuf(), pic_out.get_vbuf()) { |
| 110 | scale_loop!(sbuf, dbuf); |
| 111 | } else if let (Some(ref sbuf), Some(ref mut dbuf)) = (pic_in.get_vbuf16(), pic_out.get_vbuf16()) { |
| 112 | scale_loop!(sbuf, dbuf); |
| 113 | } else if let (Some(ref sbuf), Some(ref mut dbuf)) = (pic_in.get_vbuf32(), pic_out.get_vbuf32()) { |
| 114 | scale_loop!(sbuf, dbuf); |
| 115 | } else { |
| 116 | unreachable!(); |
| 117 | } |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | pub fn create_scale() -> Box<dyn Kernel> { |
| 122 | Box::new(NNResampler::new()) |
| 123 | } |
| 124 | |