make edge_emu() accept block alignment
[nihav.git] / nihav-codec-support / src / codecs / blockdsp.rs
1 //! Various pixel block manipulation functions.
2 use nihav_core::frame::*;
3
4 /// Puts YUV420 16x16 macroblock data onto picture in the requested place.
5 pub fn put_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[[i16;64]; 6]) {
6 let stridey = buf.get_stride(0);
7 let strideu = buf.get_stride(1);
8 let stridev = buf.get_stride(2);
9 let mut idxy = buf.get_offset(0) + xpos * 16 + ypos * 16 * stridey;
10 let mut idxu = buf.get_offset(1) + xpos * 8 + ypos * 8 * strideu;
11 let mut idxv = buf.get_offset(2) + xpos * 8 + ypos * 8 * stridev;
12
13 let data = buf.get_data_mut().unwrap();
14 let framebuf: &mut [u8] = data.as_mut_slice();
15
16 for j in 0..8 {
17 for k in 0..8 {
18 let mut v = blk[0][k + j * 8];
19 if v < 0 { v = 0; } else if v > 255 { v = 255; }
20 framebuf[idxy + k] = v as u8;
21 }
22 for k in 0..8 {
23 let mut v = blk[1][k + j * 8];
24 if v < 0 { v = 0; } else if v > 255 { v = 255; }
25 framebuf[idxy + k + 8] = v as u8;
26 }
27 idxy += stridey;
28 }
29 for j in 0..8 {
30 for k in 0..8 {
31 let mut v = blk[2][k + j * 8];
32 if v < 0 { v = 0; } else if v > 255 { v = 255; }
33 framebuf[idxy + k] = v as u8;
34 }
35 for k in 0..8 {
36 let mut v = blk[3][k + j * 8];
37 if v < 0 { v = 0; } else if v > 255 { v = 255; }
38 framebuf[idxy + k + 8] = v as u8;
39 }
40 idxy += stridey;
41 }
42
43 for j in 0..8 {
44 for k in 0..8 {
45 let mut v = blk[4][k + j * 8];
46 if v < 0 { v = 0; } else if v > 255 { v = 255; }
47 framebuf[idxu + k] = v as u8;
48 }
49 for k in 0..8 {
50 let mut v = blk[5][k + j * 8];
51 if v < 0 { v = 0; } else if v > 255 { v = 255; }
52 framebuf[idxv + k] = v as u8;
53 }
54 idxu += strideu;
55 idxv += stridev;
56 }
57 }
58
59 /// Adds YUV420 16x16 macroblock coefficients to the picture in the requested place.
60 pub fn add_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[[i16;64]; 6]) {
61 let stridey = buf.get_stride(0);
62 let strideu = buf.get_stride(1);
63 let stridev = buf.get_stride(2);
64 let mut idxy = buf.get_offset(0) + xpos * 16 + ypos * 16 * stridey;
65 let mut idxu = buf.get_offset(1) + xpos * 8 + ypos * 8 * strideu;
66 let mut idxv = buf.get_offset(2) + xpos * 8 + ypos * 8 * stridev;
67
68 let data = buf.get_data_mut().unwrap();
69 let framebuf: &mut [u8] = data.as_mut_slice();
70
71 for j in 0..8 {
72 for k in 0..8 {
73 let mut v = blk[0][k + j * 8] + i16::from(framebuf[idxy + k]);
74 if v < 0 { v = 0; } else if v > 255 { v = 255; }
75 framebuf[idxy + k] = v as u8;
76 }
77 for k in 0..8 {
78 let mut v = blk[1][k + j * 8] + i16::from(framebuf[idxy + k + 8]);
79 if v < 0 { v = 0; } else if v > 255 { v = 255; }
80 framebuf[idxy + k + 8] = v as u8;
81 }
82 idxy += stridey;
83 }
84 for j in 0..8 {
85 for k in 0..8 {
86 let mut v = blk[2][k + j * 8] + i16::from(framebuf[idxy + k]);
87 if v < 0 { v = 0; } else if v > 255 { v = 255; }
88 framebuf[idxy + k] = v as u8;
89 }
90 for k in 0..8 {
91 let mut v = blk[3][k + j * 8] + i16::from(framebuf[idxy + k + 8]);
92 if v < 0 { v = 0; } else if v > 255 { v = 255; }
93 framebuf[idxy + k + 8] = v as u8;
94 }
95 idxy += stridey;
96 }
97
98 for j in 0..8 {
99 for k in 0..8 {
100 let mut v = blk[4][k + j * 8] + i16::from(framebuf[idxu + k]);
101 if v < 0 { v = 0; } else if v > 255 { v = 255; }
102 framebuf[idxu + k] = v as u8;
103 }
104 for k in 0..8 {
105 let mut v = blk[5][k + j * 8] + i16::from(framebuf[idxv + k]);
106 if v < 0 { v = 0; } else if v > 255 { v = 255; }
107 framebuf[idxv + k] = v as u8;
108 }
109 idxu += strideu;
110 idxv += stridev;
111 }
112 }
113
114 /// Copies block from the picture with pixels beyond the picture borders being replaced with replicated edge pixels.
115 pub fn edge_emu(src: &NAVideoBuffer<u8>, xpos: isize, ypos: isize, bw: usize, bh: usize, dst: &mut [u8], dstride: usize, comp: usize, align: u8) {
116 let stride = src.get_stride(comp);
117 let offs = src.get_offset(comp);
118 let (w_, h_) = src.get_dimensions(comp);
119 let (hss, vss) = src.get_info().get_format().get_chromaton(comp).unwrap().get_subsampling();
120 let data = src.get_data();
121 let framebuf: &[u8] = data.as_slice();
122
123 let (w, h) = if align == 0 {
124 (w_, h_)
125 } else {
126 let wa = if align > hss { (1 << (align - hss)) - 1 } else { 0 };
127 let ha = if align > vss { (1 << (align - vss)) - 1 } else { 0 };
128 ((w_ + wa) & !wa, (h_ + ha) & !ha)
129 };
130
131 for y in 0..bh {
132 let srcy;
133 if (y as isize) + ypos < 0 { srcy = 0; }
134 else if (y as isize) + ypos >= (h as isize) { srcy = h - 1; }
135 else { srcy = ((y as isize) + ypos) as usize; }
136
137 for x in 0..bw {
138 let srcx;
139 if (x as isize) + xpos < 0 { srcx = 0; }
140 else if (x as isize) + xpos >= (w as isize) { srcx = w - 1; }
141 else { srcx = ((x as isize) + xpos) as usize; }
142 dst[x + y * dstride] = framebuf[offs + srcx + srcy * stride];
143 }
144 }
145 }
146
147 /// A generic type for motion interpolation function used by [`copy_blocks`]
148 ///
149 /// The function expects following parameters:
150 /// * destination buffer
151 /// * destination buffer stride
152 /// * source buffer
153 /// * source buffer stride
154 /// * block width
155 /// * block height
156 ///
157 /// [`copy_blocks`]: ./fn.copy_blocks.html
158 pub type BlkInterpFunc = fn(&mut [u8], usize, &[u8], usize, usize, usize);
159
160 /// Performs motion compensation on YUV420 macroblock.
161 ///
162 /// Arguments:
163 /// * `dx` and `dy` - destination coordinates
164 /// * `sx` and `sy` - source coordinates
165 /// * `bw` and `bh` - block dimensions
166 /// * `preborder` and `postborder` - number of pixels before and after interpolated one used by the interpolation filter.
167 /// * `mode` - interpolation mode (essentially the index for the `interp` array)
168 pub fn copy_blocks(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>,
169 dx: usize, dy: usize, sx: isize, sy: isize, bw: usize, bh: usize,
170 preborder: usize, postborder: usize,
171 mode: usize, interp: &[BlkInterpFunc])
172 {
173 let pre = if mode != 0 { preborder as isize } else { 0 };
174 let post = if mode != 0 { postborder as isize } else { 0 };
175 let (w, h) = src.get_dimensions(0);
176
177 if (sx - pre < 0) || ((sx >> 1) - pre < 0) || (sx + (bw as isize) + post > (w as isize)) ||
178 (sy - pre < 0) || ((sy >> 1) - pre < 0) || (sy + (bh as isize) + post > (h as isize)) {
179 let ebuf_stride: usize = 32;
180 let mut ebuf: Vec<u8> = vec![0; ebuf_stride * (bh + ((pre + post) as usize))];
181
182 for comp in 0..3 {
183 let dstride = dst.get_stride(comp);
184 let doff = dst.get_offset(comp);
185 let ddta = dst.get_data_mut().unwrap();
186 let dbuf: &mut [u8] = ddta.as_mut_slice();
187 let x = if comp > 0 { dx/2 } else { dx };
188 let y = if comp > 0 { dy/2 } else { dy };
189 let sx_ = (if comp > 0 { sx >> 1 } else { sx }) - pre;
190 let sy_ = (if comp > 0 { sy >> 1 } else { sy }) - pre;
191 let bw_ = (if comp > 0 { bw/2 } else { bw }) + ((pre + post) as usize);
192 let bh_ = (if comp > 0 { bh/2 } else { bh }) + ((pre + post) as usize);
193 edge_emu(src, sx_ - pre, sy_ - pre, bw_, bh_,
194 ebuf.as_mut_slice(), ebuf_stride, comp, 0);
195 let bw_ = if comp > 0 { bw/2 } else { bw };
196 let bh_ = if comp > 0 { bh/2 } else { bh };
197 (interp[mode])(&mut dbuf[doff + x + y * dstride..], dstride, ebuf.as_slice(), ebuf_stride, bw_, bh_);
198 }
199 } else {
200 for comp in 0..3 {
201 let sstride = src.get_stride(comp);
202 let soff = src.get_offset(comp);
203 let sdta = src.get_data();
204 let sbuf: &[u8] = sdta.as_slice();
205 let dstride = dst.get_stride(comp);
206 let doff = dst.get_offset(comp);
207 let ddta = dst.get_data_mut().unwrap();
208 let dbuf: &mut [u8] = ddta.as_mut_slice();
209 let x = if comp > 0 { dx/2 } else { dx };
210 let y = if comp > 0 { dy/2 } else { dy };
211 let sx_ = ((if comp > 0 { sx >> 1 } else { sx }) - pre) as usize;
212 let sy_ = ((if comp > 0 { sy >> 1 } else { sy }) - pre) as usize;
213 let bw_ = if comp > 0 { bw/2 } else { bw };
214 let bh_ = if comp > 0 { bh/2 } else { bh };
215 (interp[mode])(&mut dbuf[doff + x + y * dstride..], dstride, &sbuf[(soff + sx_ + sy_ * sstride)..], sstride, bw_, bh_);
216 }
217 }
218 }
219
220 /// Performs motion compensation on arbitrary block on some plane.
221 ///
222 /// See [`copy_blocks`] for the arguments explanation.
223 ///
224 /// [`copy_blocks`]: ./fn.copy_blocks.html
225 pub fn copy_block(dst: &mut NASimpleVideoFrame<u8>, src: NAVideoBufferRef<u8>, comp: usize,
226 dx: usize, dy: usize, mv_x: i16, mv_y: i16, bw: usize, bh: usize,
227 preborder: usize, postborder: usize,
228 mode: usize, interp: &[BlkInterpFunc])
229 {
230 let pre = if mode != 0 { preborder as isize } else { 0 };
231 let post = if mode != 0 { postborder as isize } else { 0 };
232 let (w, h) = src.get_dimensions(comp);
233 let sx = (dx as isize) + (mv_x as isize);
234 let sy = (dy as isize) + (mv_y as isize);
235
236 if (sx - pre < 0) || (sx + (bw as isize) + post > (w as isize)) ||
237 (sy - pre < 0) || (sy + (bh as isize) + post > (h as isize)) {
238 let ebuf_stride: usize = 32;
239 let mut ebuf: Vec<u8> = vec![0; ebuf_stride * (bh + ((pre + post) as usize))];
240
241 let dstride = dst.stride[comp];
242 let doff = dst.offset[comp];
243 let edge = (pre + post) as usize;
244 edge_emu(&src, sx - pre, sy - pre, bw + edge, bh + edge,
245 ebuf.as_mut_slice(), ebuf_stride, comp, 0);
246 (interp[mode])(&mut dst.data[doff + dx + dy * dstride..], dstride,
247 ebuf.as_slice(), ebuf_stride, bw, bh);
248 } else {
249 let sstride = src.get_stride(comp);
250 let soff = src.get_offset(comp);
251 let sdta = src.get_data();
252 let sbuf: &[u8] = sdta.as_slice();
253 let dstride = dst.stride[comp];
254 let doff = dst.offset[comp];
255 let saddr = soff + ((sx - pre) as usize) + ((sy - pre) as usize) * sstride;
256 (interp[mode])(&mut dst.data[doff + dx + dy * dstride..], dstride,
257 &sbuf[saddr..], sstride, bw, bh);
258 }
259 }