lrf.rs 20.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
// Copyright (c) 2017-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.

#![allow(safe_extern_statics)]

12
13
14
15
16
use encoder::Frame;
use encoder::FrameInvariants;
use context::ContextWriter;
use context::SuperBlockOffset;
use context::PLANES;
Monty's avatar
Monty committed
17
18
19
use context::MAX_SB_SIZE;
use plane::Plane;
use plane::PlaneOffset;
20
use plane::PlaneConfig;
21
use std::cmp;
Monty's avatar
Monty committed
22
use util::clamp;
23

24
pub const RESTORATION_TILESIZE_MAX_LOG2: usize = 8;
25

26
27
28
29
30
pub const RESTORE_NONE: u8 = 0;
pub const RESTORE_SWITCHABLE: u8 = 1;
pub const RESTORE_WIENER: u8 = 2;
pub const RESTORE_SGRPROJ: u8 = 3;

31
32
33
pub const WIENER_TAPS_MIN: [i8; 3] = [ -5, -23, -17 ];
pub const WIENER_TAPS_MID: [i8; 3] = [ 3, -7, 15 ];
pub const WIENER_TAPS_MAX: [i8; 3] = [ 10, 8, 46 ];
34
pub const WIENER_TAPS_K:   [i8; 3] = [ 1, 2, 3 ];
Monty's avatar
Monty committed
35
pub const WIENER_BITS: usize = 7;
36
37
38
39
40
41
42

pub const SGRPROJ_XQD_MIN: [i8; 2] = [ -96, -32 ];
pub const SGRPROJ_XQD_MID: [i8; 2] = [ -32, 31 ];
pub const SGRPROJ_XQD_MAX: [i8; 2] = [ 31, 95 ];
pub const SGRPROJ_PRJ_SUBEXP_K: u8 = 4;
pub const SGRPROJ_PRJ_BITS: u8 = 7;
pub const SGRPROJ_PARAMS_BITS: u8 = 4;
Monty's avatar
Monty committed
43
44
45
46
pub const SGRPROJ_MTABLE_BITS: u8 = 20;
pub const SGRPROJ_SGR_BITS: u8 = 8;
pub const SGRPROJ_RECIP_BITS: u8 = 12;
pub const SGRPROJ_RST_BITS: u8 = 4;
47
48
49
50
51
52
53
54
55
56
57
58
59
pub const SGRPROJ_PARAMS_RADIUS: [[u8; 2]; 1 << SGRPROJ_PARAMS_BITS] = [
  [2, 1], [2, 1], [2, 1], [2, 1],
  [2, 1], [2, 1], [2, 1], [2, 1],
  [2, 1], [2, 1], [0, 1], [0, 1],
  [0, 1], [0, 1], [2, 0], [2, 0],
];
pub const SGRPROJ_PARAMS_EPS: [[u8; 2]; 1 << SGRPROJ_PARAMS_BITS] = [
  [12,  4], [15,  6], [18,  8], [21,  9],
  [24, 10], [29, 11], [36, 12], [45, 13],
  [56, 14], [68, 15], [ 0,  5], [ 0,  8],
  [ 0, 11], [ 0, 14], [30,  0], [75,  0],
];

60
#[derive(Copy, Clone, Debug)]
61
62
63
pub enum RestorationFilter {
  None,
  Wiener  { coeffs: [[i8; 3]; 2] },
Monty's avatar
Monty committed
64
  Sgrproj { set: u8,
65
66
67
68
69
            xqd: [i8; 2] },
}

impl RestorationFilter {
  pub fn default() -> RestorationFilter {
70
    RestorationFilter::None
Monty's avatar
Monty committed
71
72
73
  }
}

Monty's avatar
Monty committed
74
75
76
fn sgrproj_box_ab(af: &mut[i32; 64+2],
                  bf: &mut[i32; 64+2],
                  r: isize, eps: isize,
77
                  fi: &FrameInvariants,
Monty's avatar
Monty committed
78
79
                  x: isize,
                  stripe_y: isize, stripe_h: isize, y: isize, h: isize,
80
81
                  cdeffed: &Plane, deblocked: &Plane) {
  let bit_depth = fi.sequence.bit_depth;
Monty's avatar
Monty committed
82
  let n = ((2*r + 1) * (2*r + 1)) as i32;
Monty's avatar
Monty committed
83
  let one_over_n = ((1 << SGRPROJ_RECIP_BITS) + n/2 ) / n;
Monty's avatar
Monty committed
84
85
  let n2e = n*n*eps as i32;
  let s = ((1 << SGRPROJ_MTABLE_BITS) + n2e/2) / n2e;
86
  let xn = cmp::min(r+1, fi.width as isize - x);
Monty's avatar
Monty committed
87
88
89
  for row in -1..1+h {
    let mut a:i32 = 0;
    let mut b:i32 = 0;
Monty's avatar
Monty committed
90

Monty's avatar
Monty committed
91
92
93
94
    for yi in y+row-r..=y+row+r {
      let mut src_plane: &Plane;
      let ly;
      if yi < stripe_y {
95
        ly = cmp::max(clamp(yi, 0, fi.height as isize - 1), y - 2) as usize;
Monty's avatar
Monty committed
96
97
        src_plane = deblocked;
      } else if yi < stripe_y + stripe_h {
98
        ly = clamp(yi, 0, fi.height as isize - 1) as usize;
Monty's avatar
Monty committed
99
        src_plane = cdeffed;
Monty's avatar
Monty committed
100
      } else {
101
        ly = cmp::min(clamp(yi, 0, fi.height as isize - 1), y+h+1) as usize;
Monty's avatar
Monty committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
        src_plane = deblocked;
      }

      for _xi in -r..-x {
        let c = src_plane.p(0, ly) as i32;
        a += c*c;
        b += c;
      }
      for xi in cmp::max(-r, -x)..xn {
        let c = src_plane.p((xi + x) as usize, ly) as i32;
        a += c*c;
        b += c;
      }
      for _xi in xn..r+1 {
116
        let c = src_plane.p(fi.width - 1, ly) as i32;
Monty's avatar
Monty committed
117
118
119
        a += c*c;
        b += c;
      }
Monty's avatar
Monty committed
120
    }
Monty's avatar
Monty committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    a = a + (1 << 2*(bit_depth-8) >> 1) >> 2*(bit_depth-8);
    let d = b + (1 << bit_depth - 8 >> 1) >> bit_depth - 8;
    let p = cmp::max(0, a*(n as i32) - d*d);
    let z = p*s + (1 << SGRPROJ_MTABLE_BITS >> 1) >> SGRPROJ_MTABLE_BITS;
    let a2 = if z >= 255 {
      256
    } else if z == 0 {
      1
    } else {
      ((z << SGRPROJ_SGR_BITS) + z/2) / (z+1)
    };
    let b2 = ((1 << SGRPROJ_SGR_BITS) - a2 ) * b * one_over_n;
    af[(row+1) as usize] = a2;
    bf[(row+1) as usize] = b2 + (1 << SGRPROJ_RECIP_BITS >> 1) >> SGRPROJ_RECIP_BITS;
Monty's avatar
Monty committed
135
  }
Monty's avatar
Monty committed
136
137
138
139
}

fn sgrproj_box_f(af: &[&[i32; 64+2]; 3], bf: &[&[i32; 64+2]; 3], f: &mut[i32; 64],
                 x: usize, y: usize, h: usize, cdeffed: &Plane, pass: usize) {
Monty's avatar
Monty committed
140
141
142

  for i in 0..h {
    let shift = if pass == 0 && (i&1) == 1 {4} else {5} + SGRPROJ_SGR_BITS - SGRPROJ_RST_BITS;
Monty's avatar
Monty committed
143
144
145
146
147
148
149
150
    let mut a = 0;
    let mut b = 0;
    for dy in 0..=2 {
      for dx in 0..=2 {
        let weight = if pass == 0 {
          if ((i+dy) & 1) == 0 {
            if dx == 1 {
              6
Monty's avatar
Monty committed
151
            } else {
Monty's avatar
Monty committed
152
              5
Monty's avatar
Monty committed
153
154
            }
          } else {
Monty's avatar
Monty committed
155
156
157
158
159
160
161
162
163
164
165
            0
          }
        } else {
          if dx == 1 || dy == 1 {
            4
          } else {
            3
          }
        };
        a += weight * af[dx][i+dy];
        b += weight * bf[dx][i+dy];
Monty's avatar
Monty committed
166
167
      }
    }
Monty's avatar
Monty committed
168
169
    let v = a * cdeffed.p(x, y+i) as i32 + b;
    f[i as usize] = v + (1 << shift >> 1) >> shift;
Monty's avatar
Monty committed
170
171
172
  }
}

173
fn sgrproj_stripe_rdu(set: u8, xqd: [i8; 2], fi: &FrameInvariants,
Monty's avatar
Monty committed
174
                      x: usize, w: usize, stripe_y: isize, stripe_h: isize,
175
                      cdeffed: &Plane, deblocked: &Plane, out: &mut Plane) {
Monty's avatar
Monty committed
176
177
178

  // unlike x, our y can be negative to start as the first stripe
  // starts off the top of the frame by 8 pixels, and can also run off the end of the frame
Monty's avatar
Monty committed
179
  let clipped_y = if stripe_y < 0 {0} else {stripe_y} as usize;
180
181
  let clipped_h = cmp::max(0, if stripe_y + stripe_h > fi.height as isize {
    fi.height - clipped_y
Monty's avatar
Monty committed
182
  } else {
Monty's avatar
Monty committed
183
184
    (stripe_y + stripe_h - clipped_y as isize) as usize
  });
Monty's avatar
Monty committed
185
186

  assert!(clipped_h <= 64);
Monty's avatar
Monty committed
187
188
189
190
191
192
193
194

  let mut a0: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
  let mut a1: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
  let mut b0: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
  let mut b1: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
  let mut f0: [i32; 64] = [0; 64];
  let mut f1: [i32; 64] = [0; 64];

Monty's avatar
Monty committed
195
196
197
198
199
  let r0: u8 = SGRPROJ_PARAMS_RADIUS[set as usize][0];
  let r1: u8 = SGRPROJ_PARAMS_RADIUS[set as usize][1];
  let eps0: u8 = SGRPROJ_PARAMS_EPS[set as usize][0];
  let eps1: u8 = SGRPROJ_PARAMS_EPS[set as usize][1];

Monty's avatar
Monty committed
200
  /* prime the intermediate arrays */
201
  if r0 > 0 {
202
    sgrproj_box_ab(&mut a0[0], &mut b0[0], r0 as isize, eps0 as isize, fi,
203
                   x as isize - 1, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
204
205
                   cdeffed, deblocked);
    sgrproj_box_ab(&mut a0[1], &mut b0[1], r0 as isize, eps0 as isize, fi,
206
                   x as isize, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
207
                   cdeffed, deblocked);
208
209
  }
  if r1 > 0 {
210
    sgrproj_box_ab(&mut a1[0], &mut b1[0], r1 as isize, eps1 as isize, fi,
211
                   x as isize - 1, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
212
213
                   cdeffed, deblocked);
    sgrproj_box_ab(&mut a1[1], &mut b1[1], r1 as isize, eps1 as isize, fi,
214
                   x as isize, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
215
                   cdeffed, deblocked);
216
217
  }
  
Monty's avatar
Monty committed
218
219
220
221
222
223
  /* iterate by column */
  let cdeffed_slice = cdeffed.slice(&PlaneOffset{x: x as isize, y: clipped_y as isize});
  let outstride = out.cfg.stride;
  let mut out_slice = out.mut_slice(&PlaneOffset{x: x as isize, y: clipped_y as isize});
  let out_data = out_slice.as_mut_slice();
  for xi in 0..w {
224
225
    /* build intermediate array columns */
    if r0 > 0 {
226
      sgrproj_box_ab(&mut a0[(xi+2)%3], &mut b0[(xi+2)%3], r0 as isize, eps0 as isize, fi,
227
                     (x + xi + 1) as isize, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
228
                     cdeffed, deblocked);
Monty's avatar
Monty committed
229
230
      let ap0: [&[i32; 64+2]; 3] = [&a0[xi%3], &a0[(xi+1)%3], &a0[(xi+2)%3]];
      let bp0: [&[i32; 64+2]; 3] = [&b0[xi%3], &b0[(xi+1)%3], &b0[(xi+2)%3]];
231
232
233
      sgrproj_box_f(&ap0, &bp0, &mut f0, x + xi, clipped_y, clipped_h, cdeffed, 0);
    }
    if r1 > 0 {
234
      sgrproj_box_ab(&mut a1[(xi+2)%3], &mut b1[(xi+2)%3], r1 as isize, eps1 as isize, fi,
235
                     (x + xi + 1) as isize, stripe_y, stripe_h, clipped_y as isize, clipped_h as isize,
236
                     cdeffed, deblocked);
237
      let ap1: [&[i32; 64+2]; 3] = [&a1[xi%3], &a1[(xi+1)%3], &a1[(xi+2)%3]];
Monty's avatar
Monty committed
238
239
240
241
      let bp1: [&[i32; 64+2]; 3] = [&b1[xi%3], &b1[(xi+1)%3], &b1[(xi+2)%3]];

      sgrproj_box_f(&ap1, &bp1, &mut f1, x + xi, clipped_y, clipped_h, cdeffed, 1);
    }
242
    let bit_depth = fi.sequence.bit_depth;
243
244
245
246
247
248
249
250
251
252
253
    if r0 > 0 {
      if r1 > 0 {
        let w0 = xqd[0] as i32;
        let w1 = xqd[1] as i32;
        let w2 = (1 << SGRPROJ_PRJ_BITS) - w0 - w1;
        for yi in 0..clipped_h {
          let u = (cdeffed_slice.p(xi,yi) as i32) << SGRPROJ_RST_BITS;
          let v = w0*f0[yi] + w1*u + w2*f1[yi];
          let s = v + (1 << SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS >> 1) >> SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS;
          out_data[xi + yi*outstride] = clamp(s as u16, 0, (1 << bit_depth) - 1);
        }
Monty's avatar
Monty committed
254
      } else {
255
256
257
258
259
260
261
262
        let w0 = xqd[0] as i32;
        let w = (1 << SGRPROJ_PRJ_BITS) - w0;
        for yi in 0..clipped_h {
          let u = (cdeffed_slice.p(xi,yi) as i32) << SGRPROJ_RST_BITS;
          let v = w0*f0[yi] + w*u;
          let s = v + (1 << SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS >> 1) >> SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS;
          out_data[xi + yi*outstride] = clamp(s as u16, 0, (1 << bit_depth) - 1);
        }
Monty's avatar
Monty committed
263
      }
264
265
266
267
268
269
270
271
272
273
    } else {
      /* if r0 is 0, the r1 must be nonzero */
      let w = xqd[0] as i32 + xqd[1] as i32;
      let w2 = (1 << SGRPROJ_PRJ_BITS) - w;
      for yi in 0..clipped_h {
        let u = (cdeffed_slice.p(xi,yi) as i32) << SGRPROJ_RST_BITS;
        let v = w*u + w2*f1[yi];
        let s = v + (1 << SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS >> 1) >> SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS;
        out_data[xi + yi*outstride] = clamp(s as u16, 0, (1 << bit_depth) - 1);
      }      
Monty's avatar
Monty committed
274
275
276
277
    }
  }
}

278
fn wiener_stripe_rdu(coeffs: [[i8; 3]; 2], fi: &FrameInvariants,
Monty's avatar
Monty committed
279
                     x: usize, w: usize, y: isize, h: isize,
280
281
                     cdeffed: &Plane, deblocked: &Plane, out: &mut Plane) {
  let bit_depth = fi.sequence.bit_depth;
Monty's avatar
Monty committed
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
  let round_h = if bit_depth == 12 {5} else {3};
  let round_v = if bit_depth == 12 {9} else {11};
  let offset = 1 << bit_depth + WIENER_BITS - round_h - 1;
  let limit = (1 << bit_depth + 1 + WIENER_BITS - round_h) - 1;
  
  let mut work: [i32; MAX_SB_SIZE+7] = [0; MAX_SB_SIZE+7];
  let vfilter: [i32; 7] = [ coeffs[0][0] as i32,
                            coeffs[0][1] as i32,
                            coeffs[0][2] as i32,
                            128 - 2 * (coeffs[0][0] as i32 +
                                       coeffs[0][1] as i32 +
                                       coeffs[0][2] as i32 ),
                            coeffs[0][2] as i32,
                            coeffs[0][1] as i32,
                            coeffs[0][0] as i32];
  let hfilter: [i32; 7] = [ coeffs[1][0] as i32,
                            coeffs[1][1] as i32,
                            coeffs[1][2] as i32,
                            128 - 2 * (coeffs[1][0] as i32 +
                                       coeffs[1][1] as i32 +
                                       coeffs[1][2] as i32),
                            coeffs[1][2] as i32,
                            coeffs[1][1] as i32,
                            coeffs[1][0] as i32];

  // unlike x, our y can be negative to start as the first stripe
  // starts off the top of the frame by 8 pixels, and can also run off the end of the frame
  let start_wi = if y < 0 {-y} else {0} as usize; 
  let start_yi = if y < 0 {0} else {y} as usize; 
311
312
  let end_i = cmp::max(0, if y+h > fi.height as isize {
    fi.height as isize - y - start_wi as isize
Monty's avatar
Monty committed
313
314
315
316
317
318
319
320
321
  } else {
    h - start_wi as isize
  }) as usize;
  
  let stride = out.cfg.stride;
  let mut out_slice = out.mut_slice(&PlaneOffset{x: 0, y: start_yi as isize});
  let out_data = out_slice.as_mut_slice();

  for xi in x..x+w {
322
    let n = cmp::min(7, fi.width as isize + 3 - xi as isize);
Monty's avatar
Monty committed
323
324
325
326
327
    for yi in y-3..y+h+4 {
      let mut src_plane: &Plane;
      let mut acc = 0;
      let ly;
      if yi < y {
328
        ly = cmp::max(clamp(yi, 0, fi.height as isize - 1), y - 2) as usize;
Monty's avatar
Monty committed
329
330
        src_plane = deblocked;
      } else if yi < y+h {
331
        ly = clamp(yi, 0, fi.height as isize - 1) as usize;
Monty's avatar
Monty committed
332
333
        src_plane = cdeffed;
      } else {
334
        ly = cmp::min(clamp(yi, 0, fi.height as isize - 1), y+h+1) as usize;
Monty's avatar
Monty committed
335
336
337
338
339
340
341
342
343
344
        src_plane = deblocked;
      }
      
      for i in 0..3 - xi as isize {
        acc += hfilter[i as usize] * src_plane.p(0, ly) as i32;
      }
      for i in cmp::max(0,3 - (xi as isize))..n {
        acc += hfilter[i as usize] * src_plane.p((xi as isize + i - 3) as usize, ly) as i32;
      }
      for i in n..7 {
345
        acc += hfilter[i as usize] * src_plane.p(fi.width - 1, ly) as i32;
Monty's avatar
Monty committed
346
347
348
349
350
351
352
353
354
355
356
357
358
      }
        
      acc = acc + (1 << round_h >> 1) >> round_h;
      work[(yi-y+3) as usize] = clamp(acc, -offset, limit-offset);
    }

    for (wi, dst) in (start_wi..start_wi+end_i).zip(out_data[xi..].iter_mut().step_by(stride).take(end_i)) {
      let mut acc = 0;
      for (i,src) in (0..7).zip(work[wi..wi+7].iter_mut()) {
        acc += vfilter[i] * *src;
      }
      *dst = clamp(acc + (1 << round_v >> 1) >> round_v, 0, (1 << bit_depth) - 1) as u16;
    }
359
360
361
  }
}

362
#[derive(Copy, Clone, Debug)]
363
364
365
366
367
368
369
370
371
372
373
374
375
376
pub struct RestorationUnit {
  pub filter: RestorationFilter,
  pub coded: bool,
}

impl RestorationUnit {
  pub fn default() -> RestorationUnit {
    RestorationUnit {
      filter: RestorationFilter::default(),
      coded: false,
    }
  }
}

377
#[derive(Clone, Debug)]
378
379
380
pub struct RestorationPlane {
  pub lrf_type: u8,
  pub unit_size: usize,
381
  // (1 << sb_shift) gives the number of superblocks having size 1 << SUPERBLOCK_TO_PLANE_SHIFT
382
  // both horizontally and vertically in a restoration unit, not accounting for RU stretching
383
  pub sb_shift: usize,
Monty's avatar
Monty committed
384
385
386
  // stripe height is 64 in all cases except 4:2:0 chroma planes where
  // it is 32.  This is independent of all other setup parameters
  pub stripe_height: usize,
387
388
389
390
  pub cols: usize,
  pub rows: usize,
  pub wiener_ref: [[i8; 3]; 2],
  pub sgrproj_ref: [i8; 2],
391
  pub units: Box<[RestorationUnit]>,
392
393
394
395
396
397
398
399
400
}

#[derive(Clone, Default)]
pub struct RestorationPlaneOffset {
  pub row: usize,
  pub col: usize
}

impl RestorationPlane {
Monty's avatar
Monty committed
401
  pub fn new(lrf_type: u8, unit_size: usize, sb_shift: usize, stripe_decimate: usize,
402
             cols: usize, rows: usize) -> RestorationPlane {
Monty's avatar
Monty committed
403
    let stripe_height = if stripe_decimate != 0 {32} else {64};
404
405
406
    RestorationPlane {
      lrf_type,
      unit_size,
407
      sb_shift,
Monty's avatar
Monty committed
408
      stripe_height,
409
410
411
412
      cols,
      rows,
      wiener_ref: [WIENER_TAPS_MID; 2],
      sgrproj_ref: SGRPROJ_XQD_MID,
413
      units: vec![RestorationUnit::default(); cols * rows].into_boxed_slice(),
414
415
416
    }
  }

417
418
  fn restoration_unit_index(&self, sbo: &SuperBlockOffset) -> (usize, usize) {
    (
419
420
      (sbo.x >> self.sb_shift).min(self.cols - 1),
      (sbo.y >> self.sb_shift).min(self.rows - 1),
421
    )
422
423
  }

424
425
426
  // Stripes are always 64 pixels high in a non-subsampled
  // frame, and decimated from 64 pixels in chroma.  When
  // filtering, they are not co-located on Y with superblocks.
Monty's avatar
Monty committed
427
  fn restoration_unit_index_by_stripe(&self, stripenum: usize, rux: usize) -> (usize, usize) {
428
429
    (
      cmp::min(rux, self.cols - 1),
Monty's avatar
Monty committed
430
      cmp::min(stripenum * self.stripe_height / self.unit_size, self.rows - 1),
431
432
433
    )
  }

434
  pub fn restoration_unit(&self, sbo: &SuperBlockOffset) -> &RestorationUnit {
435
    let (x, y) = self.restoration_unit_index(sbo);
436
    &self.units[y * self.cols + x]
437
438
439
  }

  pub fn restoration_unit_as_mut(&mut self, sbo: &SuperBlockOffset) -> &mut RestorationUnit {
440
    let (x, y) = self.restoration_unit_index(sbo);
441
    &mut self.units[y * self.cols + x]
Monty's avatar
Monty committed
442
  }
443

Monty's avatar
Monty committed
444
445
  pub fn restoration_unit_by_stripe(&self, stripenum: usize, rux: usize) -> &RestorationUnit {
    let (x, y) = self.restoration_unit_index_by_stripe(stripenum, rux);
446
447
    &self.units[y * self.cols + x]
  }
448
449
}

450
#[derive(Clone, Debug)]
451
452
453
454
455
pub struct RestorationState {
  pub plane: [RestorationPlane; PLANES]
}

impl RestorationState {
456
  pub fn new(fi: &FrameInvariants, input: &Frame) -> Self {
457
    let PlaneConfig { xdec, ydec, .. } = input.planes[1].cfg;
Monty's avatar
Monty committed
458
    let stripe_uv_decimate = if xdec>0 && ydec>0 {1} else {0};
459
460
461
462
463
464
    // Currrently opt for smallest possible restoration unit size (1
    // superblock) This is *temporary*.  Counting on it will break
    // very shortly; the 1-superblock hardwiring is only until the
    // upper level encoder is capable of dealing with the delayed
    // writes that RU size > SB size will require.
    let lrf_y_shift = if fi.sequence.use_128x128_superblock {1} else {2};
Monty's avatar
Monty committed
465
    let lrf_uv_shift = lrf_y_shift + stripe_uv_decimate;
466
467
468
469
470
471
472

    // derive the rest
    let y_unit_log2 = RESTORATION_TILESIZE_MAX_LOG2 - lrf_y_shift;
    let uv_unit_log2 = RESTORATION_TILESIZE_MAX_LOG2 - lrf_uv_shift;
    let y_unit_size = 1 << y_unit_log2;
    let uv_unit_size = 1 << uv_unit_log2;
    let y_sb_log2 = if fi.sequence.use_128x128_superblock {7} else {6};
Monty's avatar
Monty committed
473
    let uv_sb_log2 = y_sb_log2 - stripe_uv_decimate;
474
475
    let cols = ((fi.width + (y_unit_size >> 1)) / y_unit_size).max(1);
    let rows = ((fi.height + (y_unit_size >> 1)) / y_unit_size).max(1);
476

477
    RestorationState {
Monty's avatar
Monty committed
478
479
480
481
482
483
      plane: [RestorationPlane::new(RESTORE_SWITCHABLE, y_unit_size, y_unit_log2 - y_sb_log2,
                                    0, cols, rows),
              RestorationPlane::new(RESTORE_SWITCHABLE, uv_unit_size, uv_unit_log2 - uv_sb_log2,
                                    stripe_uv_decimate, cols, rows),
              RestorationPlane::new(RESTORE_SWITCHABLE, uv_unit_size, uv_unit_log2 - uv_sb_log2,
                                    stripe_uv_decimate, cols, rows)],
484
485
    }
  }
486

487
  pub fn restoration_unit(&self, sbo: &SuperBlockOffset, pli: usize) -> &RestorationUnit {
488
    self.plane[pli].restoration_unit(sbo)
489
490
491
  }

  pub fn restoration_unit_as_mut(&mut self, sbo: &SuperBlockOffset, pli: usize) -> &mut RestorationUnit {
492
493
    self.plane[pli].restoration_unit_as_mut(sbo)
  }
494
495

  pub fn lrf_optimize_superblock(&mut self, _sbo: &SuperBlockOffset, _fi: &FrameInvariants,
496
                                 _cw: &mut ContextWriter) {
497
498
  }

499
  pub fn lrf_filter_frame(&mut self, out: &mut Frame, pre_cdef: &Frame,
500
                          fi: &FrameInvariants) {
501
    let cdeffed = out.clone();
Monty's avatar
Monty committed
502
    
503
504
505
506
507
    // unlike the other loop filters that operate over the padded
    // frame dimensions, restoration filtering and source pixel
    // accesses are clipped to the original frame dimensions
    // that's why we use fi.width and fi.height instead of PlaneConfig fields

Monty's avatar
Monty committed
508
    // number of stripes (counted according to colocated Y luma position)
509
    let stripe_n = (fi.height + 7) / 64 + 1;
Monty's avatar
Monty committed
510
511
512
    
    for pli in 0..PLANES {
      let rp = &self.plane[pli];
513
      let ydec = out.planes[pli].cfg.ydec;
Monty's avatar
Monty committed
514
515
516
517
518
519
520
521
522
523
524

      for si in 0..stripe_n {
        // stripe y pixel locations must be able to overspan the frame
        let stripe_start_y = si as isize * 64 - 8 >> ydec;
        let stripe_size = 64 >> ydec; // one past, unlike spec
      
        // horizontally, go rdu-by-rdu
        for rux in 0..rp.cols {
          // stripe x pixel locations must be clipped to frame, last may need to stretch
          let ru_start_x = rux * rp.unit_size;
          let ru_size = if rux == rp.cols - 1 {
525
            fi.width - ru_start_x
Monty's avatar
Monty committed
526
527
528
          } else {
            rp.unit_size
          };
Monty's avatar
Monty committed
529
          let ru = rp.restoration_unit_by_stripe(si, rux);
Monty's avatar
Monty committed
530
531
          match ru.filter {
            RestorationFilter::Wiener{coeffs} => {          
532
              wiener_stripe_rdu(coeffs, fi,
Monty's avatar
Monty committed
533
534
                                ru_start_x, ru_size, stripe_start_y, stripe_size,
                                &cdeffed.planes[pli], &pre_cdef.planes[pli],
535
                                &mut out.planes[pli]);
Monty's avatar
Monty committed
536
537
            },
            RestorationFilter::Sgrproj{set, xqd} => {
538
              sgrproj_stripe_rdu(set, xqd, fi,
Monty's avatar
Monty committed
539
540
                                 ru_start_x, ru_size, stripe_start_y, stripe_size,
                                 &cdeffed.planes[pli], &pre_cdef.planes[pli],
541
                                 &mut out.planes[pli]);
Monty's avatar
Monty committed
542
543
544
545
546
547
548
549
            },
            RestorationFilter::None => {
              // do nothing
            }
          }
        }
      }
    }    
550
551
  }
}