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>

4 files changed, 93 insertions, 130 deletions

diff --git a/libavcodec/aacenc.c b/libavcodec/aacenc.c
index 232eeda04f..1a845be407 100644
--- a/libavcodec/aacenc.c
+++ b/libavcodec/aacenc.c
@@ -611,7 +611,7 @@ static int aac_encode_frame(AVCodecContext *avctx, AVPacket *avpkt,
                 if (s->options.tns && s->coder->search_for_tns)
                     s->coder->search_for_tns(s, sce);
                 if (s->options.tns && s->coder->apply_tns_filt)
-                    s->coder->apply_tns_filt(sce);
+                    s->coder->apply_tns_filt(s, sce);
                 if (sce->tns.present)
                     tns_mode = 1;
             }
diff --git a/libavcodec/aacenc.h b/libavcodec/aacenc.h
index 51dce8a0fd..2b7a62a146 100644
--- a/libavcodec/aacenc.h
+++ b/libavcodec/aacenc.h
@@ -63,7 +63,7 @@ typedef struct AACCoefficientsEncoder {
     void (*encode_main_pred)(struct AACEncContext *s, SingleChannelElement *sce);
     void (*adjust_common_prediction)(struct AACEncContext *s, ChannelElement *cpe);
     void (*apply_main_pred)(struct AACEncContext *s, SingleChannelElement *sce);
-    void (*apply_tns_filt)(SingleChannelElement *sce);
+    void (*apply_tns_filt)(struct AACEncContext *s, SingleChannelElement *sce);
     void (*set_special_band_scalefactors)(struct AACEncContext *s, SingleChannelElement *sce);
     void (*search_for_pns)(struct AACEncContext *s, AVCodecContext *avctx, SingleChannelElement *sce);
     void (*search_for_tns)(struct AACEncContext *s, SingleChannelElement *sce);
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++;
         }
diff --git a/libavcodec/aacenc_tns.h b/libavcodec/aacenc_tns.h
index 72c91239a3..812deeac8c 100644
--- a/libavcodec/aacenc_tns.h
+++ b/libavcodec/aacenc_tns.h
@@ -30,11 +30,21 @@
 
 #include "aacenc.h"
 
-/** Lower limit of TNS in SFBs **/
-#define TNS_LOW_LIMIT 24
+/* Could be set to 3 to save an additional bit at the cost of little quality */
+#define TNS_Q_BITS 4
+
+/* TNS will only be used if the LPC gain is within these margins */
+#define TNS_GAIN_THRESHOLD_LOW  1.395f
+#define TNS_GAIN_THRESHOLD_HIGH 11.19f
+
+/* Do not use TNS if the psy band spread is below this value */
+#define TNS_SPREAD_THRESHOLD 20.081512f
+
+/* Allows to reverse the filter direction if the band energy is uneven */
+#define TNS_DIRECTION_VARY 1
 
 void ff_aac_encode_tns_info(AACEncContext *s, SingleChannelElement *sce);
-void ff_aac_apply_tns(SingleChannelElement *sce);
+void ff_aac_apply_tns(AACEncContext *s, SingleChannelElement *sce);
 void ff_aac_search_for_tns(AACEncContext *s, SingleChannelElement *sce);
 
 #endif /* AVCODEC_AACENC_TNS_H */