aacenc_tns: rework coefficient quantization and filter application

This commit reworks the TNS implementation to a hybrid between what
the specifications say, what the decoder does and what's the best
thing to do.

The filter application function was copied from the decoder and
modified such that it applies the inverse AR filter to the
coefficients. The LPC coefficients themselves are fed into the
same quantization expression that the specifications say should
be used however further processing is not done, instead they're
converted to the form that the decoder expects them to be in
and are sent off to the compute_lpc_coeffs function exactly the
way the decoder does. This function does all conversions and will
return the exact coefficients that the decoder will generate, which
are then applied to the coefficients.
Having the exact same coefficients on both the encoder and decoder
is a must since otherwise the entire sfb's over which the filter
is applied will be attenuated.

Despite this major rework, TNS might not work fine on some audio
types at very low bitrates (e.g. sub 90kbps) as it can attenuate
some coefficients too much. Users are advised to experiment with
TNS at higher bitrates if they wish to use this tool or simply
wait for the implementation to be improved.

Signed-off-by: Rostislav Pehlivanov <atomnuker@gmail.com>

1 files changed, 78 insertions, 125 deletions

diff --git a/libavcodec/aacenc_tns.c b/libavcodec/aacenc_tns.c
index d4d10e68ff..3c442e81d6 100644
--- a/libavcodec/aacenc_tns.c
+++ b/libavcodec/aacenc_tns.c
@@ -31,112 +31,80 @@
 #include "aacenc_utils.h"
 #include "aacenc_quantization.h"
 
-static inline int compress_coef(int *coefs, int num)
-{
-    int i, c = 0;
-    for (i = 0; i < num; i++)
-        c += coefs[i] < 4 || coefs[i] > 11;
-    return c == num;
-}
-
 /**
  * Encode TNS data.
  * Coefficient compression saves a single bit per coefficient.
  */
 void ff_aac_encode_tns_info(AACEncContext *s, SingleChannelElement *sce)
 {
-    int i, w, filt, coef_len, coef_compress;
+    uint8_t u_coef;
+    const uint8_t coef_res = TNS_Q_BITS == 4;
+    int i, w, filt, coef_len, coef_compress = 0;
     const int is8 = sce->ics.window_sequence[0] == EIGHT_SHORT_SEQUENCE;
+    TemporalNoiseShaping *tns = &sce->tns;
 
     if (!sce->tns.present)
         return;
 
     for (i = 0; i < sce->ics.num_windows; i++) {
         put_bits(&s->pb, 2 - is8, sce->tns.n_filt[i]);
-        if (sce->tns.n_filt[i]) {
-            put_bits(&s->pb, 1, 1);
-            for (filt = 0; filt < sce->tns.n_filt[i]; filt++) {
-                put_bits(&s->pb, 6 - 2 * is8, sce->tns.length[i][filt]);
-                put_bits(&s->pb, 5 - 2 * is8, sce->tns.order[i][filt]);
-                if (sce->tns.order[i][filt]) {
-                    coef_compress = compress_coef(sce->tns.coef_idx[i][filt],
-                                                  sce->tns.order[i][filt]);
-                    put_bits(&s->pb, 1, !!sce->tns.direction[i][filt]);
+        if (tns->n_filt[i]) {
+            put_bits(&s->pb, 1, coef_res);
+            for (filt = 0; filt < tns->n_filt[i]; filt++) {
+                put_bits(&s->pb, 6 - 2 * is8, tns->length[i][filt]);
+                put_bits(&s->pb, 5 - 2 * is8, tns->order[i][filt]);
+                if (tns->order[i][filt]) {
+                    put_bits(&s->pb, 1, !!tns->direction[i][filt]);
                     put_bits(&s->pb, 1, !!coef_compress);
-                    coef_len = 4 - coef_compress;
-                    for (w = 0; w < sce->tns.order[i][filt]; w++)
-                        put_bits(&s->pb, coef_len, sce->tns.coef_idx[i][filt][w]);
+                    coef_len = coef_res + 3 - coef_compress;
+                    for (w = 0; w < tns->order[i][filt]; w++) {
+                        u_coef = (tns->coef_idx[i][filt][w])&(~(~0<<coef_len));
+                        put_bits(&s->pb, coef_len, u_coef);
+                    }
                 }
             }
         }
     }
 }
 
-static void process_tns_coeffs(TemporalNoiseShaping *tns, double *coef_raw,
-                               int *order_p, int w, int filt)
+static int quantize_coefs(double *coef, int *idx, float *lpc, int order)
 {
-    int i, j, order = *order_p;
-    int *idx = tns->coef_idx[w][filt];
-    float *lpc = tns->coef[w][filt];
-    float temp[TNS_MAX_ORDER] = {0.0f}, out[TNS_MAX_ORDER] = {0.0f};
-
-    if (!order)
-        return;
-
-    /* Not what the specs say, but it's better */
+    int i;
+    uint8_t u_coef;
+    const float *quant_arr = tns_tmp2_map[TNS_Q_BITS == 4];
+    const double iqfac_p = ((1 << (TNS_Q_BITS-1)) - 0.5)/(M_PI/2.0);
+    const double iqfac_m = ((1 << (TNS_Q_BITS-1)) + 0.5)/(M_PI/2.0);
     for (i = 0; i < order; i++) {
-        idx[i] = quant_array_idx(coef_raw[i], tns_tmp2_map_0_4, 16);
-        lpc[i] = tns_tmp2_map_0_4[idx[i]];
-    }
-
-    /* Trim any coeff less than 0.1f from the end */
-    for (i = order-1; i > -1; i--) {
-        lpc[i] = (fabs(lpc[i]) > 0.1f) ? lpc[i] : 0.0f;
-        if (lpc[i] != 0.0 ) {
-            order = i;
-            break;
-        }
-    }
-    order = av_clip(order, 0, TNS_MAX_ORDER - 1);
-    *order_p = order;
-    if (!order)
-        return;
-
-    /* Step up procedure, convert to LPC coeffs */
-    out[0] = 1.0f;
-    for (i = 1; i <= order; i++) {
-        for (j = 1; j < i; j++) {
-            temp[j] = out[j] + lpc[i]*out[i-j];
-        }
-        for (j = 1; j <= i; j++) {
-            out[j] = temp[j];
-        }
-        out[i] = lpc[i-1];
+        idx[i] = ceilf(asin(coef[i])*((coef[i] >= 0) ? iqfac_p : iqfac_m));
+        u_coef = (idx[i])&(~(~0<<TNS_Q_BITS));
+        lpc[i] = quant_arr[u_coef];
     }
-    memcpy(lpc, out, TNS_MAX_ORDER*sizeof(float));
+    return order;
 }
 
 /* Apply TNS filter */
-void ff_aac_apply_tns(SingleChannelElement *sce)
+void ff_aac_apply_tns(AACEncContext *s, SingleChannelElement *sce)
 {
-    float *coef = sce->pcoeffs;
     TemporalNoiseShaping *tns = &sce->tns;
-    int w, filt, m, i;
-    int bottom, top, order, start, end, size, inc;
-    float *lpc, tmp[TNS_MAX_ORDER+1];
+    IndividualChannelStream *ics = &sce->ics;
+    int w, filt, m, i, top, order, bottom, start, end, size, inc;
+    const int mmm = FFMIN(ics->tns_max_bands, ics->max_sfb);
+    float lpc[TNS_MAX_ORDER];
 
-    for (w = 0; w < sce->ics.num_windows; w++) {
-        bottom = sce->ics.num_swb;
+    for (w = 0; w < ics->num_windows; w++) {
+        bottom = ics->num_swb;
         for (filt = 0; filt < tns->n_filt[w]; filt++) {
             top    = bottom;
             bottom = FFMAX(0, top - tns->length[w][filt]);
             order  = tns->order[w][filt];
-            lpc    = tns->coef[w][filt];
-            if (!order)
+            if (order == 0)
                 continue;
 
-            start = sce->ics.swb_offset[bottom];
-            end   = sce->ics.swb_offset[top];
+            // tns_decode_coef
+            compute_lpc_coefs(tns->coef[w][filt], order, lpc, 0, 0, 0);
+
+            start = ics->swb_offset[FFMIN(bottom, mmm)];
+            end   = ics->swb_offset[FFMIN(   top, mmm)];
             if ((size = end - start) <= 0)
                 continue;
             if (tns->direction[w][filt]) {
@@ -147,21 +115,10 @@ void ff_aac_apply_tns(SingleChannelElement *sce)
             }
             start += w * 128;
 
-            if (!sce->ics.ltp.present) {
-                // ar filter
-                for (m = 0; m < size; m++, start += inc)
-                    for (i = 1; i <= FFMIN(m, order); i++)
-                        coef[start] += coef[start - i * inc]*lpc[i - 1];
-            } else {
-                // ma filter
-                for (m = 0; m < size; m++, start += inc) {
-                    tmp[0] = coef[start];
-                    for (i = 1; i <= FFMIN(m, order); i++)
-                        coef[start] += tmp[i]*lpc[i - 1];
-                    for (i = order; i > 0; i--)
-                        tmp[i] = tmp[i - 1];
-                }
-            }
+            // ar filter
+            for (m = 0; m < size; m++, start += inc)
+                for (i = 1; i <= FFMIN(m, order); i++)
+                    sce->coeffs[start] += lpc[i-1]*sce->pcoeffs[start - i*inc];
         }
     }
 }
@@ -169,57 +126,53 @@ void ff_aac_apply_tns(SingleChannelElement *sce)
 void ff_aac_search_for_tns(AACEncContext *s, SingleChannelElement *sce)
 {
     TemporalNoiseShaping *tns = &sce->tns;
-    int w, g, w2, prev_end_sfb = 0, count = 0;
+    int w, w2, g, count = 0;
+    const int mmm = FFMIN(sce->ics.tns_max_bands, sce->ics.max_sfb);
     const int is8 = sce->ics.window_sequence[0] == EIGHT_SHORT_SEQUENCE;
-    const int tns_max_order = is8 ? 7 : s->profile == FF_PROFILE_AAC_LOW ? 12 : TNS_MAX_ORDER;
+    int order = is8 ? 7 : s->profile == FF_PROFILE_AAC_LOW ? 12 : TNS_MAX_ORDER;
+
+    int sfb_start = av_clip(tns_min_sfb[is8][s->samplerate_index], 0, mmm);
+    int sfb_end   = av_clip(sce->ics.num_swb, 0, mmm);
 
     for (w = 0; w < sce->ics.num_windows; w++) {
-        int order = 0, filters = 1;
-        int sfb_start = 0, sfb_len = 0;
-        int coef_start = 0, coef_len = 0;
-        float energy = 0.0f, threshold = 0.0f;
-        double coefs[MAX_LPC_ORDER][MAX_LPC_ORDER] = {{0}};
+        float en_low = 0.0f, en_high = 0.0f, threshold = 0.0f, spread = 0.0f;
+        double gain = 0.0f, coefs[MAX_LPC_ORDER] = {0};
+
+        int coef_start = w*sce->ics.num_swb + sce->ics.swb_offset[sfb_start];
+        int coef_len = sce->ics.swb_offset[sfb_end] - sce->ics.swb_offset[sfb_start];
+
         for (g = 0;  g < sce->ics.num_swb; g++) {
-            if (!sfb_start && w*16+g > TNS_LOW_LIMIT && w*16+g > prev_end_sfb) {
-                sfb_start = w*16+g;
-                coef_start =  sce->ics.swb_offset[sfb_start];
-            }
-            if (sfb_start) {
-                for (w2 = 0; w2 < sce->ics.group_len[w]; w2++) {
-                    FFPsyBand *band = &s->psy.ch[s->cur_channel].psy_bands[(w+w2)*16+g];
-                    if (!sfb_len && band->energy < band->threshold*1.3f) {
-                        sfb_len = (w+w2)*16+g - sfb_start;
-                        prev_end_sfb = sfb_start + sfb_len;
-                        coef_len = sce->ics.swb_offset[sfb_start + sfb_len] - coef_start;
-                        break;
-                    }
-                    energy += band->energy;
-                    threshold += band->threshold;
-                }
-                if (!sfb_len) {
-                    sfb_len = (w+1)*16+g - sfb_start - 1;
-                    coef_len = sce->ics.swb_offset[sfb_start + sfb_len] - coef_start;
-                }
+            if (w*16+g < sfb_start || w*16+g > sfb_end)
+                continue;
+            for (w2 = 0; w2 < sce->ics.group_len[w]; w2++) {
+                FFPsyBand *band = &s->psy.ch[s->cur_channel].psy_bands[(w+w2)*16+g];
+                if ((w+w2)*16+g > sfb_start + ((sfb_end - sfb_start)/2))
+                    en_high += band->energy;
+                else
+                    en_low  += band->energy;
+                threshold += band->threshold;
+                spread += band->spread;
             }
         }
 
-        if (sfb_len <= 0 || coef_len <= 0)
+        if (coef_len <= 0 || (sfb_end - sfb_start) <= 0)
             continue;
-        if (coef_start + coef_len >= 1024)
-            coef_len = 1024 - coef_start;
 
         /* LPC */
-        order = ff_lpc_calc_levinson(&s->lpc, &sce->coeffs[coef_start], coef_len,
-                                     coefs, 0, tns_max_order, ORDER_METHOD_LOG);
+        gain = ff_lpc_calc_ref_coefs_f(&s->lpc, &sce->coeffs[coef_start],
+                                       coef_len, order, coefs);
+
+        gain *= s->lambda/110.0f;
 
-        if (energy > threshold) {
-            int direction = 0;
-            tns->n_filt[w] = filters++;
+        if (gain > TNS_GAIN_THRESHOLD_LOW && gain*0 < TNS_GAIN_THRESHOLD_HIGH &&
+            (en_low+en_high) > TNS_GAIN_THRESHOLD_LOW*threshold &&
+            spread > TNS_SPREAD_THRESHOLD) {
+            tns->n_filt[w] = 1;
             for (g = 0; g < tns->n_filt[w]; g++) {
-                process_tns_coeffs(tns, coefs[order], &order, w, g);
-                tns->order[w][g]     = order;
-                tns->length[w][g]    = sfb_len;
-                tns->direction[w][g] = direction;
+                tns->length[w][g] = sfb_end - sfb_start;
+                tns->direction[w][g] = en_low < en_high && TNS_DIRECTION_VARY;
+                tns->order[w][g] = quantize_coefs(coefs, tns->coef_idx[w][g],
+                                                  tns->coef[w][g], order);
             }
             count++;
         }