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 }
 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