test_encode_decode.rs 8.79 KB
Newer Older
rzumer's avatar
rzumer committed
1 2 3 4 5 6 7 8 9
// Copyright (c) 2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

10 11 12 13 14 15 16 17 18 19 20
#![allow(dead_code)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(non_upper_case_globals)]

include!(concat!(env!("OUT_DIR"), "/aom.rs"));

use super::*;

use rand::{ChaChaRng, Rng, SeedableRng};
use std::collections::VecDeque;
21
use std::mem;
Luca Barbato's avatar
Luca Barbato committed
22
use std::sync::Arc;
23

24 25 26 27 28 29 30 31
fn fill_frame(ra: &mut ChaChaRng, frame: &mut Frame) {
  for plane in frame.planes.iter_mut() {
    let stride = plane.cfg.stride;
    for row in plane.data.chunks_mut(stride) {
      for mut pixel in row {
        let v: u8 = ra.gen();
        *pixel = v as u16;
      }
32
    }
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
  }
}

struct AomDecoder {
  dec: aom_codec_ctx
}

fn setup_decoder(w: usize, h: usize) -> AomDecoder {
  unsafe {
    let interface = aom_codec_av1_dx();
    let mut dec: AomDecoder = mem::uninitialized();
    let cfg = aom_codec_dec_cfg_t {
      threads: 1,
      w: w as u32,
      h: h as u32,
      allow_lowbitdepth: 1,
      cfg: cfg_options { ext_partition: 1 }
    };

    let ret = aom_codec_dec_init_ver(
      &mut dec.dec,
      interface,
      &cfg,
      0,
      AOM_DECODER_ABI_VERSION as i32
    );
    if ret != 0 {
      panic!("Cannot instantiate the decoder {}", ret);
61 62
    }

63 64 65 66 67 68 69 70 71 72 73 74
    dec
  }
}

impl Drop for AomDecoder {
  fn drop(&mut self) {
    unsafe { aom_codec_destroy(&mut self.dec) };
  }
}

fn setup_encoder(
  w: usize, h: usize, speed: usize, quantizer: usize, bit_depth: usize,
Josh Holmer's avatar
Josh Holmer committed
75 76
  chroma_sampling: ChromaSampling, min_keyint: u64, max_keyint: u64,
  low_latency: bool
77 78 79 80 81 82
) -> Context {
  unsafe {
    av1_rtcd();
    aom_dsp_rtcd();
  }

83 84
  let mut enc = EncoderConfig::with_speed_preset(speed);
  enc.quantizer = quantizer;
Josh Holmer's avatar
Josh Holmer committed
85 86 87
  enc.min_key_frame_interval = min_keyint;
  enc.max_key_frame_interval = max_keyint;
  enc.low_latency = low_latency;
88 89

  let cfg = Config {
90 91 92 93 94 95 96 97
    frame_info: FrameInfo { 
      width: w, 
      height: h, 
      bit_depth, 
      chroma_sampling, 
      chroma_sample_position: 
      Default::default() 
    },
Luca Barbato's avatar
Luca Barbato committed
98
    timebase: Rational::new(1, 1000),
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118
    enc
  };

  cfg.new_context()
}

// TODO: support non-multiple-of-16 dimensions
static DIMENSION_OFFSETS: &[(usize, usize)] =
  &[(0, 0), (4, 4), (8, 8), (16, 16)];

#[test]
#[ignore]
fn speed() {
  let quantizer = 100;
  let limit = 5;
  let w = 64;
  let h = 80;

  for b in DIMENSION_OFFSETS.iter() {
    for s in 0..10 {
Josh Holmer's avatar
Josh Holmer committed
119
      encode_decode(w + b.0, h + b.1, s, quantizer, limit, 8, 15, 15, true);
120
    }
121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148
  }
}

static DIMENSIONS: &[(usize, usize)] = &[
  (8, 8),
  (16, 16),
  (32, 32),
  (64, 64),
  (128, 128),
  (256, 256),
  (512, 512),
  (1024, 1024),
  (2048, 2048),
  (258, 258),
  (260, 260),
  (262, 262),
  (264, 264),
  (265, 265)
];

#[test]
#[ignore]
fn dimensions() {
  let quantizer = 100;
  let limit = 1;
  let speed = 4;

  for (w, h) in DIMENSIONS.iter() {
Josh Holmer's avatar
Josh Holmer committed
149
    encode_decode(*w, *h, speed, quantizer, limit, 8, 15, 15, true);
150 151 152 153 154 155 156 157 158 159 160 161
  }
}

#[test]
fn quantizer() {
  let limit = 5;
  let w = 64;
  let h = 80;
  let speed = 4;

  for b in DIMENSION_OFFSETS.iter() {
    for &q in [80, 100, 120].iter() {
Josh Holmer's avatar
Josh Holmer committed
162
      encode_decode(w + b.0, h + b.1, speed, q, limit, 8, 15, 15, true);
163
    }
164 165 166
  }
}

Josh Holmer's avatar
Josh Holmer committed
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190
#[test]
fn keyframes() {
  let limit = 12;
  let w = 64;
  let h = 80;
  let speed = 10;
  let q = 100;

  encode_decode(w, h, speed, q, limit, 8, 6, 6, true);
}

#[test]
fn reordering() {
  let limit = 12;
  let w = 64;
  let h = 80;
  let speed = 10;
  let q = 100;

  for keyint in &[4, 5, 6] {
    encode_decode(w, h, speed, q, limit, 8, *keyint, *keyint, false);
  }
}

191 192 193 194 195 196 197 198 199
#[test]
#[ignore]
fn odd_size_frame_with_full_rdo() {
  let limit = 3;
  let w = 512 + 32 + 16 + 5;
  let h = 512 + 16 + 5;
  let speed = 0;
  let qindex = 100;

Josh Holmer's avatar
Josh Holmer committed
200
  encode_decode(w, h, speed, qindex, limit, 8, 15, 15, true);
201 202 203 204 205 206 207 208 209 210 211
}

#[test]
fn high_bd() {
  let quantizer = 100;
  let limit = 3; // Include inter frames
  let speed = 0; // Test as many tools as possible
  let w = 64;
  let h = 80;

  // 10-bit
Josh Holmer's avatar
Josh Holmer committed
212
  encode_decode(w, h, speed, quantizer, limit, 10, 15, 15, true);
213 214

  // 12-bit
Josh Holmer's avatar
Josh Holmer committed
215
  encode_decode(w, h, speed, quantizer, limit, 12, 15, 15, true);
216 217 218 219 220 221 222 223 224 225 226
}

fn compare_plane<T: Ord + std::fmt::Debug>(
  rec: &[T], rec_stride: usize, dec: &[T], dec_stride: usize, width: usize,
  height: usize
) {
  for line in rec.chunks(rec_stride).zip(dec.chunks(dec_stride)).take(height) {
    assert_eq!(&line.0[..width], &line.1[..width]);
  }
}

Frank Bossen's avatar
Frank Bossen committed
227
fn compare_img(img: *const aom_image_t, frame: &Frame, bit_depth: usize, width: usize, height: usize) {
228 229 230 231 232
  use std::slice;
  let img = unsafe { *img };
  let img_iter = img.planes.iter().zip(img.stride.iter());

  for (img_plane, frame_plane) in img_iter.zip(frame.planes.iter()) {
Frank Bossen's avatar
Frank Bossen committed
233 234
    let w = width >> frame_plane.cfg.xdec;
    let h = height >> frame_plane.cfg.ydec;
235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
    let rec_stride = frame_plane.cfg.stride;

    if bit_depth > 8 {
      let dec_stride = *img_plane.1 as usize / 2;

      let dec = unsafe {
        let data = *img_plane.0 as *const u16;
        let size = dec_stride * h;

        slice::from_raw_parts(data, size)
      };

      let rec: Vec<u16> =
        frame_plane.data_origin().iter().map(|&v| v).collect();

      compare_plane::<u16>(&rec[..], rec_stride, dec, dec_stride, w, h);
    } else {
      let dec_stride = *img_plane.1 as usize;

      let dec = unsafe {
        let data = *img_plane.0 as *const u8;
        let size = dec_stride * h;

        slice::from_raw_parts(data, size)
      };

      let rec: Vec<u8> =
        frame_plane.data_origin().iter().map(|&v| v as u8).collect();

      compare_plane::<u8>(&rec[..], rec_stride, dec, dec_stride, w, h);
265
    }
266 267 268 269 270
  }
}

fn encode_decode(
  w: usize, h: usize, speed: usize, quantizer: usize, limit: usize,
Josh Holmer's avatar
Josh Holmer committed
271
  bit_depth: usize, min_keyint: u64, max_keyint: u64, low_latency: bool
272 273 274 275 276 277
) {
  use std::ptr;
  let mut ra = ChaChaRng::from_seed([0; 32]);

  let mut dec = setup_decoder(w, h);
  let mut ctx =
Josh Holmer's avatar
Josh Holmer committed
278 279
    setup_encoder(w, h, speed, quantizer, bit_depth, ChromaSampling::Cs420,
                  min_keyint, max_keyint, low_latency);
280 281 282 283 284 285 286 287 288 289 290 291 292

  println!("Encoding {}x{} speed {} quantizer {}", w, h, speed, quantizer);

  let mut iter: aom_codec_iter_t = ptr::null_mut();

  let mut rec_fifo = VecDeque::new();

  for _ in 0..limit {
    let mut input = ctx.new_frame();
    fill_frame(&mut ra, Arc::get_mut(&mut input).unwrap());

    let _ = ctx.send_frame(input);

fbossen's avatar
fbossen committed
293
    let mut done = false;
294
    let mut corrupted_count = 0;
fbossen's avatar
fbossen committed
295 296 297 298 299 300 301
    while !done {
      let res = ctx.receive_packet();
      if let Ok(pkt) = res {
        println!("Encoded packet {}", pkt.number);

        if let Some(pkt_rec) = pkt.rec {
          rec_fifo.push_back(pkt_rec.clone());
302 303
        }

fbossen's avatar
fbossen committed
304
        let packet = pkt.data;
305

fbossen's avatar
fbossen committed
306 307 308
        unsafe {
          println!("Decoding frame {}", pkt.number);
          let ret = aom_codec_decode(
309
            &mut dec.dec,
fbossen's avatar
fbossen committed
310 311 312
            packet.as_ptr(),
            packet.len(),
            ptr::null_mut()
313
          );
fbossen's avatar
fbossen committed
314
          println!("Decoded. -> {}", ret);
315 316
          if ret != 0 {
            use std::ffi::CStr;
fbossen's avatar
fbossen committed
317 318 319 320
            let error_msg = aom_codec_error(&mut dec.dec);
            println!(
              "  Decode codec_decode failed: {}",
              CStr::from_ptr(error_msg).to_string_lossy()
321
            );
fbossen's avatar
fbossen committed
322 323 324 325 326 327 328 329 330
            let detail = aom_codec_error_detail(&mut dec.dec);
            if !detail.is_null() {
              println!(
                "  Decode codec_decode failed {}",
                CStr::from_ptr(detail).to_string_lossy()
              );
            }

            corrupted_count += 1;
331 332
          }

fbossen's avatar
fbossen committed
333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357
          if ret == 0 {
            loop {
              println!("Retrieving frame");
              let img = aom_codec_get_frame(&mut dec.dec, &mut iter);
              println!("Retrieved.");
              if img.is_null() {
                break;
              }
              let mut corrupted = 0;
              let ret = aom_codec_control_(
                &mut dec.dec,
                aom_dec_control_id_AOMD_GET_FRAME_CORRUPTED as i32,
                &mut corrupted
              );
              if ret != 0 {
                use std::ffi::CStr;
                let detail = aom_codec_error_detail(&mut dec.dec);
                panic!(
                  "Decode codec_control failed {}",
                  CStr::from_ptr(detail).to_string_lossy()
                );
              }
              corrupted_count += corrupted;

              let rec = rec_fifo.pop_front().unwrap();
Frank Bossen's avatar
Frank Bossen committed
358
              compare_img(img, &rec, bit_depth, w, h);
fbossen's avatar
fbossen committed
359 360
            }
          }
361
        }
fbossen's avatar
fbossen committed
362 363
      } else {
        done = true;
364
      }
365
    }
366 367 368
    assert_eq!(corrupted_count, 0);
  }
}