lrf.rs 20.7 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)]

Raphaël Zumer's avatar
Raphaël Zumer committed
12
13
14
15
16
17
18
19
20
use crate::encoder::Frame;
use crate::encoder::FrameInvariants;
use crate::context::ContextWriter;
use crate::context::SuperBlockOffset;
use crate::context::PLANES;
use crate::context::MAX_SB_SIZE;
use crate::plane::Plane;
use crate::plane::PlaneOffset;
use crate::plane::PlaneConfig;
21
use std::cmp;
Raphaël Zumer's avatar
Raphaël Zumer committed
22
use crate::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,
Monty's avatar
Monty committed
77
78
79
                  crop_w: usize, crop_h: usize, stripe_h: isize,
                  stripe_x: isize, stripe_y: isize,
                  cdeffed: &Plane, deblocked: &Plane, bit_depth: usize) {
Monty's avatar
Monty committed
80
  let n = ((2*r + 1) * (2*r + 1)) as i32;
Monty's avatar
Monty committed
81
  let one_over_n = ((1 << SGRPROJ_RECIP_BITS) + n/2 ) / n;
Monty's avatar
Monty committed
82
83
  let n2e = n*n*eps as i32;
  let s = ((1 << SGRPROJ_MTABLE_BITS) + n2e/2) / n2e;
Monty's avatar
Monty committed
84
85
  let xn = cmp::min(r+1, crop_w as isize - stripe_x);
  for row in -1..1+stripe_h {
Monty's avatar
Monty committed
86
87
    let mut a:i32 = 0;
    let mut b:i32 = 0;
Monty's avatar
Monty committed
88

Monty's avatar
Monty committed
89
    for yi in stripe_y+row-r..=stripe_y+row+r {
Raphaël Zumer's avatar
Raphaël Zumer committed
90
      let src_plane: &Plane;
Monty's avatar
Monty committed
91
92
      let ly;
      if yi < stripe_y {
Monty's avatar
Monty committed
93
        ly = cmp::max(cmp::max(yi, 0), stripe_y - 2) as usize;
Monty's avatar
Monty committed
94
95
        src_plane = deblocked;
      } else if yi < stripe_y + stripe_h {
Monty's avatar
Monty committed
96
        ly = clamp(yi, 0, crop_h as isize - 1) as usize;
Monty's avatar
Monty committed
97
        src_plane = cdeffed;
Monty's avatar
Monty committed
98
      } else {
Monty's avatar
Monty committed
99
        ly = cmp::min(clamp(yi, 0, crop_h as isize - 1), stripe_y+stripe_h+1) as usize;
Monty's avatar
Monty committed
100
101
102
        src_plane = deblocked;
      }

Monty's avatar
Monty committed
103
      for _xi in -r..-stripe_x {
Monty's avatar
Monty committed
104
105
106
107
        let c = src_plane.p(0, ly) as i32;
        a += c*c;
        b += c;
      }
Monty's avatar
Monty committed
108
109
      for xi in cmp::max(-r, -stripe_x)..xn {
        let c = src_plane.p((xi + stripe_x) as usize, ly) as i32;
Monty's avatar
Monty committed
110
111
112
113
        a += c*c;
        b += c;
      }
      for _xi in xn..r+1 {
Monty's avatar
Monty committed
114
        let c = src_plane.p(crop_w - 1, ly) as i32;
Monty's avatar
Monty committed
115
116
117
        a += c*c;
        b += c;
      }
Monty's avatar
Monty committed
118
    }
Monty's avatar
Monty committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
    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
133
  }
Monty's avatar
Monty committed
134
135
136
}

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

Monty's avatar
Monty committed
139
  for i in 0..h as usize {
Monty's avatar
Monty committed
140
    let shift = if pass == 0 && (i&1) == 1 {4} else {5} + SGRPROJ_SGR_BITS - SGRPROJ_RST_BITS;
Monty's avatar
Monty committed
141
142
143
144
145
146
147
148
    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
149
            } else {
Monty's avatar
Monty committed
150
              5
Monty's avatar
Monty committed
151
152
            }
          } else {
Monty's avatar
Monty committed
153
154
155
156
157
158
159
160
161
162
163
            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
164
165
      }
    }
Monty's avatar
Monty committed
166
    let v = a * cdeffed.p(x, (y+i as isize) as usize) as i32 + b;
Monty's avatar
Monty committed
167
    f[i as usize] = v + (1 << shift >> 1) >> shift;
Monty's avatar
Monty committed
168
169
170
  }
}

171
fn sgrproj_stripe_rdu(set: u8, xqd: [i8; 2], fi: &FrameInvariants,
Monty's avatar
Monty committed
172
173
174
                      crop_w: usize, crop_h: usize,
                      stripe_w: usize, stripe_h: isize,
                      stripe_x: usize, stripe_y: isize,
175
                      cdeffed: &Plane, deblocked: &Plane, out: &mut Plane) {
Monty's avatar
Monty committed
176

Monty's avatar
Monty committed
177
  assert!(stripe_h <= 64);
Monty's avatar
Monty committed
178
179
180
181
182
183
184
  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
185
186
187
188
189
  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
190
  /* prime the intermediate arrays */
191
  if r0 > 0 {
Monty's avatar
Monty committed
192
193
194
195
196
197
198
199
    sgrproj_box_ab(&mut a0[0], &mut b0[0], r0 as isize, eps0 as isize,
                   crop_w, crop_h, stripe_h,
                   stripe_x as isize - 1, stripe_y,
                   cdeffed, deblocked, fi.sequence.bit_depth);
    sgrproj_box_ab(&mut a0[1], &mut b0[1], r0 as isize, eps0 as isize,
                   crop_w, crop_h, stripe_h,
                   stripe_x as isize, stripe_y,
                   cdeffed, deblocked, fi.sequence.bit_depth);
200
201
  }
  if r1 > 0 {
Monty's avatar
Monty committed
202
203
204
205
206
207
208
209
    sgrproj_box_ab(&mut a1[0], &mut b1[0], r1 as isize, eps1 as isize,
                   crop_w, crop_h, stripe_h,
                   stripe_x as isize - 1, stripe_y,
                   cdeffed, deblocked, fi.sequence.bit_depth);
    sgrproj_box_ab(&mut a1[1], &mut b1[1], r1 as isize, eps1 as isize,
                   crop_w, crop_h, stripe_h,
                   stripe_x as isize, stripe_y,
                   cdeffed, deblocked, fi.sequence.bit_depth);
210
  }
Raphaël Zumer's avatar
Raphaël Zumer committed
211

Monty's avatar
Monty committed
212
  /* iterate by column */
Monty's avatar
Monty committed
213
214
  let start = cmp::max(0, -stripe_y) as usize;
  let cdeffed_slice = cdeffed.slice(&PlaneOffset{x: stripe_x as isize, y: stripe_y});
Raphaël Zumer's avatar
Raphaël Zumer committed
215
  let outstride = out.cfg.stride;
Monty's avatar
Monty committed
216
  let mut out_slice = out.mut_slice(&PlaneOffset{x: stripe_x as isize, y: stripe_y});
Monty's avatar
Monty committed
217
  let out_data = out_slice.as_mut_slice();
Monty's avatar
Monty committed
218
  for xi in 0..stripe_w {
219
220
    /* build intermediate array columns */
    if r0 > 0 {
Monty's avatar
Monty committed
221
222
223
224
      sgrproj_box_ab(&mut a0[(xi+2)%3], &mut b0[(xi+2)%3], r0 as isize, eps0 as isize,
                     crop_w, crop_h, stripe_h,
                     (stripe_x + xi + 1) as isize, stripe_y,
                     cdeffed, deblocked, fi.sequence.bit_depth);
Monty's avatar
Monty committed
225
226
      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]];
Monty's avatar
Monty committed
227
      sgrproj_box_f(&ap0, &bp0, &mut f0, stripe_x + xi, stripe_y, stripe_h, cdeffed, 0);
228
229
    }
    if r1 > 0 {
Monty's avatar
Monty committed
230
231
232
233
      sgrproj_box_ab(&mut a1[(xi+2)%3], &mut b1[(xi+2)%3], r1 as isize, eps1 as isize,
                     crop_w, crop_h, stripe_h,
                     (stripe_x + xi + 1) as isize, stripe_y,
                     cdeffed, deblocked, fi.sequence.bit_depth);
234
      let ap1: [&[i32; 64+2]; 3] = [&a1[xi%3], &a1[(xi+1)%3], &a1[(xi+2)%3]];
Monty's avatar
Monty committed
235
236
      let bp1: [&[i32; 64+2]; 3] = [&b1[xi%3], &b1[(xi+1)%3], &b1[(xi+2)%3]];

Monty's avatar
Monty committed
237
      sgrproj_box_f(&ap1, &bp1, &mut f1, stripe_x + xi, stripe_y, stripe_h, cdeffed, 1);
Monty's avatar
Monty committed
238
    }
239
    let bit_depth = fi.sequence.bit_depth;
240
241
242
243
244
    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;
Monty's avatar
Monty committed
245
246
        for yi in start..stripe_h as usize {
          let u = (cdeffed_slice.p(xi, yi) as i32) << SGRPROJ_RST_BITS;
247
248
249
250
          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
251
      } else {
252
253
        let w0 = xqd[0] as i32;
        let w = (1 << SGRPROJ_PRJ_BITS) - w0;
Monty's avatar
Monty committed
254
255
        for yi in start..stripe_h as usize {
          let u = (cdeffed_slice.p(xi, yi) as i32) << SGRPROJ_RST_BITS;
256
257
258
259
          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
260
      }
261
262
263
264
    } 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;
Monty's avatar
Monty committed
265
266
      for yi in start..stripe_h as usize {
        let u = (cdeffed_slice.p(xi, yi) as i32) << SGRPROJ_RST_BITS;
267
268
269
        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);
Raphaël Zumer's avatar
Raphaël Zumer committed
270
      }
Monty's avatar
Monty committed
271
272
273
274
    }
  }
}

275
fn wiener_stripe_rdu(coeffs: [[i8; 3]; 2], fi: &FrameInvariants,
Monty's avatar
Monty committed
276
277
278
279
                      crop_w: usize, crop_h: usize,
                      stripe_w: usize, stripe_h: isize,
                      stripe_x: usize, stripe_y: isize,
                      cdeffed: &Plane, deblocked: &Plane, out: &mut Plane) {
280
  let bit_depth = fi.sequence.bit_depth;
Monty's avatar
Monty committed
281
282
283
284
  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;
Raphaël Zumer's avatar
Raphaël Zumer committed
285

Monty's avatar
Monty committed
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
  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
Monty's avatar
Monty committed
308
309
310
311
  let start_wi = if stripe_y < 0 {-stripe_y} else {0} as usize;
  let start_yi = if stripe_y < 0 {0} else {stripe_y} as usize;
  let end_i = cmp::max(0, if stripe_y + stripe_h > crop_h as isize {
    crop_h as isize - stripe_y - start_wi as isize
Monty's avatar
Monty committed
312
  } else {
Monty's avatar
Monty committed
313
    stripe_h - start_wi as isize
Monty's avatar
Monty committed
314
  }) as usize;
Raphaël Zumer's avatar
Raphaël Zumer committed
315

Monty's avatar
Monty committed
316
317
318
319
  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();

Monty's avatar
Monty committed
320
321
322
  for xi in stripe_x..stripe_x+stripe_w {
    let n = cmp::min(7, crop_w as isize + 3 - xi as isize);
    for yi in stripe_y-3..stripe_y+stripe_h+4 {
Raphaël Zumer's avatar
Raphaël Zumer committed
323
      let src_plane: &Plane;
Monty's avatar
Monty committed
324
325
      let mut acc = 0;
      let ly;
Monty's avatar
Monty committed
326
327
      if yi < stripe_y {
        ly = cmp::max(clamp(yi, 0, crop_h as isize - 1), stripe_y - 2) as usize;
Monty's avatar
Monty committed
328
        src_plane = deblocked;
Monty's avatar
Monty committed
329
330
      } else if yi < stripe_y+stripe_h {
        ly = clamp(yi, 0, crop_h as isize - 1) as usize;
Monty's avatar
Monty committed
331
332
        src_plane = cdeffed;
      } else {
Monty's avatar
Monty committed
333
        ly = cmp::min(clamp(yi, 0, crop_h as isize - 1), stripe_y + stripe_h + 1) as usize;
Monty's avatar
Monty committed
334
335
        src_plane = deblocked;
      }
Raphaël Zumer's avatar
Raphaël Zumer committed
336

Monty's avatar
Monty committed
337
338
339
340
341
342
343
      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 {
Monty's avatar
Monty committed
344
        acc += hfilter[i as usize] * src_plane.p(crop_w - 1, ly) as i32;
Monty's avatar
Monty committed
345
      }
Raphaël Zumer's avatar
Raphaël Zumer committed
346

Monty's avatar
Monty committed
347
      acc = acc + (1 << round_h >> 1) >> round_h;
Monty's avatar
Monty committed
348
      work[(yi-stripe_y+3) as usize] = clamp(acc, -offset, limit-offset);
Monty's avatar
Monty committed
349
350
351
352
353
354
355
356
357
    }

    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;
    }
358
359
360
  }
}

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

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

376
#[derive(Clone, Debug)]
377
378
379
pub struct RestorationPlane {
  pub lrf_type: u8,
  pub unit_size: usize,
380
  // (1 << sb_shift) gives the number of superblocks having size 1 << SUPERBLOCK_TO_PLANE_SHIFT
381
  // both horizontally and vertically in a restoration unit, not accounting for RU stretching
382
  pub sb_shift: usize,
Monty's avatar
Monty committed
383
384
385
  // 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,
386
387
388
389
  pub cols: usize,
  pub rows: usize,
  pub wiener_ref: [[i8; 3]; 2],
  pub sgrproj_ref: [i8; 2],
390
  pub units: Box<[RestorationUnit]>,
391
392
393
394
395
396
397
398
399
}

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

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

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

423
424
425
  // 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
426
  fn restoration_unit_index_by_stripe(&self, stripenum: usize, rux: usize) -> (usize, usize) {
427
428
    (
      cmp::min(rux, self.cols - 1),
Monty's avatar
Monty committed
429
      cmp::min(stripenum * self.stripe_height / self.unit_size, self.rows - 1),
430
431
432
    )
  }

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

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

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

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

impl RestorationState {
455
  pub fn new(fi: &FrameInvariants, input: &Frame) -> Self {
456
    let PlaneConfig { xdec, ydec, .. } = input.planes[1].cfg;
Monty's avatar
Monty committed
457
    let stripe_uv_decimate = if xdec>0 && ydec>0 {1} else {0};
458
459
460
461
462
463
    // 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
464
    let lrf_uv_shift = lrf_y_shift + stripe_uv_decimate;
465
466
467
468
469
470
471

    // 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
472
    let uv_sb_log2 = y_sb_log2 - stripe_uv_decimate;
473
474
    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);
475

476
    RestorationState {
Monty's avatar
Monty committed
477
478
479
480
481
482
      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)],
483
484
    }
  }
485

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

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

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

498
  pub fn lrf_filter_frame(&mut self, out: &mut Frame, pre_cdef: &Frame,
499
                          fi: &FrameInvariants) {
500
    let cdeffed = out.clone();
Raphaël Zumer's avatar
Raphaël Zumer committed
501

502
503
504
505
506
    // 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
507
    // number of stripes (counted according to colocated Y luma position)
508
    let stripe_n = (fi.height + 7) / 64 + 1;
Raphaël Zumer's avatar
Raphaël Zumer committed
509

Monty's avatar
Monty committed
510
511
    for pli in 0..PLANES {
      let rp = &self.plane[pli];
Monty's avatar
Monty committed
512
      let xdec = out.planes[pli].cfg.xdec;
513
      let ydec = out.planes[pli].cfg.ydec;
Monty's avatar
Monty committed
514
515
      let crop_w = fi.width + (1 << xdec >> 1) >> xdec;
      let crop_h = fi.height + (1 << ydec >> 1) >> ydec;
Raphaël Zumer's avatar
Raphaël Zumer committed
516

Monty's avatar
Monty committed
517
518
519
520
      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
Raphaël Zumer's avatar
Raphaël Zumer committed
521

Monty's avatar
Monty committed
522
523
524
525
526
        // 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 {
Monty's avatar
Monty committed
527
            crop_w - ru_start_x
Monty's avatar
Monty committed
528
529
530
          } else {
            rp.unit_size
          };
Monty's avatar
Monty committed
531
          let ru = rp.restoration_unit_by_stripe(si, rux);
Monty's avatar
Monty committed
532
          match ru.filter {
Raphaël Zumer's avatar
Raphaël Zumer committed
533
            RestorationFilter::Wiener{coeffs} => {
534
              wiener_stripe_rdu(coeffs, fi,
Monty's avatar
Monty committed
535
536
537
                                crop_w, crop_h,
                                ru_size, stripe_size,
                                ru_start_x, stripe_start_y,
Monty's avatar
Monty committed
538
                                &cdeffed.planes[pli], &pre_cdef.planes[pli],
539
                                &mut out.planes[pli]);
Monty's avatar
Monty committed
540
541
            },
            RestorationFilter::Sgrproj{set, xqd} => {
542
              sgrproj_stripe_rdu(set, xqd, fi,
Monty's avatar
Monty committed
543
544
545
                                 crop_w, crop_h,
                                 ru_size, stripe_size,
                                 ru_start_x, stripe_start_y,
Monty's avatar
Monty committed
546
                                 &cdeffed.planes[pli], &pre_cdef.planes[pli],
547
                                 &mut out.planes[pli]);
Monty's avatar
Monty committed
548
549
550
551
552
553
554
            },
            RestorationFilter::None => {
              // do nothing
            }
          }
        }
      }
Raphaël Zumer's avatar
Raphaël Zumer committed
555
    }
556
557
  }
}