| 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 | } |
| 260 | |
| 261 | fn hpel_interp00(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize) |
| 262 | { |
| 263 | let mut didx = 0; |
| 264 | let mut sidx = 0; |
| 265 | for _ in 0..bh { |
| 266 | dst[didx..][..bw].copy_from_slice(&src[sidx..][..bw]); |
| 267 | didx += dstride; |
| 268 | sidx += sstride; |
| 269 | } |
| 270 | } |
| 271 | |
| 272 | fn hpel_interp01(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize) |
| 273 | { |
| 274 | let mut didx = 0; |
| 275 | let mut sidx = 0; |
| 276 | for _ in 0..bh { |
| 277 | for x in 0..bw { dst[didx + x] = ((u16::from(src[sidx + x]) + u16::from(src[sidx + x + 1]) + 1) >> 1) as u8; } |
| 278 | didx += dstride; |
| 279 | sidx += sstride; |
| 280 | } |
| 281 | } |
| 282 | |
| 283 | fn hpel_interp10(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize) |
| 284 | { |
| 285 | let mut didx = 0; |
| 286 | let mut sidx = 0; |
| 287 | for _ in 0..bh { |
| 288 | for x in 0..bw { dst[didx + x] = ((u16::from(src[sidx + x]) + u16::from(src[sidx + x + sstride]) + 1) >> 1) as u8; } |
| 289 | didx += dstride; |
| 290 | sidx += sstride; |
| 291 | } |
| 292 | } |
| 293 | |
| 294 | fn hpel_interp11(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize) |
| 295 | { |
| 296 | let mut didx = 0; |
| 297 | let mut sidx = 0; |
| 298 | for _ in 0..bh { |
| 299 | for x in 0..bw { |
| 300 | dst[didx + x] = ((u16::from(src[sidx + x]) + |
| 301 | u16::from(src[sidx + x + 1]) + |
| 302 | u16::from(src[sidx + x + sstride]) + |
| 303 | u16::from(src[sidx + x + sstride + 1]) + 2) >> 2) as u8; |
| 304 | } |
| 305 | didx += dstride; |
| 306 | sidx += sstride; |
| 307 | } |
| 308 | } |
| 309 | |
| 310 | /// Half-pixel interpolation functions. |
| 311 | pub const HALFPEL_INTERP_FUNCS: &[BlkInterpFunc] = &[ |
| 312 | hpel_interp00, hpel_interp01, hpel_interp10, hpel_interp11 ]; |
| 313 | |