Surround: Better LFE handling

Forces CELT-only mode for LFE (despite the rate) and "locks" most of the
CELT analysis:
- No transient or TF
- Band boost on first band
- Only first two bands get PVQ bits
- Forced energy decay after the first two bands

6 files changed, 48 insertions, 16 deletions

diff --git a/celt/celt.h b/celt/celt.h
index a1c2c480..a8f7cb03 100644
--- a/celt/celt.h
+++ b/celt/celt.h
@@ -107,7 +107,7 @@ typedef struct {
 #define CELT_SET_ANALYSIS_REQUEST    10022
 #define CELT_SET_ANALYSIS(x) CELT_SET_ANALYSIS_REQUEST, __celt_check_analysis_ptr(x)
 
-#define OPUS_SET_LFE_REQUEST    10022
+#define OPUS_SET_LFE_REQUEST    10024
 #define OPUS_SET_LFE(x) OPUS_SET_LFE_REQUEST, __opus_check_int(x)
 
 /* Encoder stuff */
diff --git a/celt/celt_encoder.c b/celt/celt_encoder.c
index 7347cb31..a88e5922 100644
--- a/celt/celt_encoder.c
+++ b/celt/celt_encoder.c
@@ -74,6 +74,7 @@ struct OpusCustomEncoder {
    int loss_rate;
    int lsb_depth;
    int variable_duration;
+   int lfe;
 
    /* Everything beyond this point gets cleared on a reset */
 #define ENCODER_RESET_START rng
@@ -869,7 +870,7 @@ static int stereo_analysis(const CELTMode *m, const celt_norm *X,
 static opus_val16 dynalloc_analysis(const opus_val16 *bandLogE, const opus_val16 *bandLogE2,
       int nbEBands, int start, int end, int C, int *offsets, int lsb_depth, const opus_int16 *logN,
       int isTransient, int vbr, int constrained_vbr, const opus_int16 *eBands, int LM,
-      int effectiveBytes, opus_int32 *tot_boost_)
+      int effectiveBytes, opus_int32 *tot_boost_, int lfe)
 {
    int i, c;
    opus_int32 tot_boost=0;
@@ -897,7 +898,7 @@ static opus_val16 dynalloc_analysis(const opus_val16 *bandLogE, const opus_val16
          maxDepth = MAX16(maxDepth, bandLogE[c*nbEBands+i]-noise_floor[i]);
    } while (++c<C);
    /* Make sure that dynamic allocation can't make us bust the budget */
-   if (effectiveBytes > 50 && LM>=1)
+   if (effectiveBytes > 50 && LM>=1 && !lfe)
    {
       int last=0;
       c=0;do
@@ -1356,7 +1357,7 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
 
    isTransient = 0;
    shortBlocks = 0;
-   if (st->complexity >= 1)
+   if (st->complexity >= 1 && !st->lfe)
    {
       isTransient = transient_analysis(in, N+st->overlap, CC,
             &tf_estimate, &tf_chan);
@@ -1429,7 +1430,7 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
 
    ALLOC(tf_res, nbEBands, int);
    /* Disable variable tf resolution for hybrid and at very low bitrate */
-   if (effectiveBytes>=15*C && st->start==0 && st->complexity>=2)
+   if (effectiveBytes>=15*C && st->start==0 && st->complexity>=2 && !st->lfe)
    {
       int lambda;
       if (effectiveBytes<40)
@@ -1455,7 +1456,7 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
    quant_coarse_energy(mode, st->start, st->end, effEnd, bandLogE,
          oldBandE, total_bits, error, enc,
          C, LM, nbAvailableBytes, st->force_intra,
-         &st->delayedIntra, st->complexity >= 4, st->loss_rate);
+         &st->delayedIntra, st->complexity >= 4, st->loss_rate, st->lfe);
 
    tf_encode(st->start, st->end, isTransient, tf_res, LM, tf_select, enc);
 
@@ -1494,7 +1495,10 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
 
    maxDepth = dynalloc_analysis(bandLogE, bandLogE2, nbEBands, st->start, st->end, C, offsets,
          st->lsb_depth, mode->logN, isTransient, st->vbr, st->constrained_vbr,
-         eBands, LM, effectiveBytes, &tot_boost);
+         eBands, LM, effectiveBytes, &tot_boost, st->lfe);
+   /* For LFE, everything interesting is in the first band */
+   if (st->lfe)
+      offsets[0] = IMIN(8, effectiveBytes/3);
    ALLOC(cap, nbEBands, int);
    init_caps(mode,cap,LM,C);
 
@@ -1560,7 +1564,10 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
    alloc_trim = 5;
    if (tell+(6<<BITRES) <= total_bits - total_boost)
    {
-      alloc_trim = alloc_trim_analysis(mode, X, bandLogE,
+      if (st->lfe)
+         alloc_trim = 5;
+      else
+         alloc_trim = alloc_trim_analysis(mode, X, bandLogE,
             st->end, LM, C, N, &st->analysis, &st->stereo_saving, tf_estimate, st->intensity);
       ec_enc_icdf(enc, alloc_trim, trim_icdf, 7);
       tell = ec_tell_frac(enc);
@@ -1738,6 +1745,8 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_val16 * pcm,
    if (st->analysis.valid)
       signalBandwidth = st->analysis.bandwidth;
 #endif
+   if (st->lfe)
+      signalBandwidth = 1;
    codedBands = compute_allocation(mode, st->start, st->end, offsets, cap,
          alloc_trim, &st->intensity, &dual_stereo, bits, &balance, pulses,
          fine_quant, fine_priority, C, LM, enc, 1, st->lastCodedBands, signalBandwidth);
@@ -2127,6 +2136,12 @@ int opus_custom_encoder_ctl(CELTEncoder * OPUS_RESTRICT st, int request, ...)
          *value=st->rng;
       }
       break;
+      case OPUS_SET_LFE_REQUEST:
+      {
+          opus_int32 value = va_arg(ap, opus_int32);
+          st->lfe = value;
+      }
+      break;
       default:
          goto bad_request;
    }
diff --git a/celt/quant_bands.c b/celt/quant_bands.c
index 514f03c4..48196bde 100644
--- a/celt/quant_bands.c
+++ b/celt/quant_bands.c
@@ -157,7 +157,7 @@ static int quant_coarse_energy_impl(const CELTMode *m, int start, int end,
       const opus_val16 *eBands, opus_val16 *oldEBands,
       opus_int32 budget, opus_int32 tell,
       const unsigned char *prob_model, opus_val16 *error, ec_enc *enc,
-      int C, int LM, int intra, opus_val16 max_decay)
+      int C, int LM, int intra, opus_val16 max_decay, int lfe)
 {
    int i, c;
    int badness = 0;
@@ -222,6 +222,8 @@ static int quant_coarse_energy_impl(const CELTMode *m, int start, int end,
             if (bits_left < 16)
                qi = IMAX(-1, qi);
          }
+         if (lfe && i>=2)
+            qi = IMIN(qi, 0);
          if (budget-tell >= 15)
          {
             int pi;
@@ -253,13 +255,13 @@ static int quant_coarse_energy_impl(const CELTMode *m, int start, int end,
          prev[c] = prev[c] + SHL32(q,7) - MULT16_16(beta,PSHR32(q,8));
       } while (++c < C);
    }
-   return badness;
+   return lfe ? 0 : badness;
 }
 
 void quant_coarse_energy(const CELTMode *m, int start, int end, int effEnd,
       const opus_val16 *eBands, opus_val16 *oldEBands, opus_uint32 budget,
       opus_val16 *error, ec_enc *enc, int C, int LM, int nbAvailableBytes,
-      int force_intra, opus_val32 *delayedIntra, int two_pass, int loss_rate)
+      int force_intra, opus_val32 *delayedIntra, int two_pass, int loss_rate, int lfe)
 {
    int intra;
    opus_val16 max_decay;
@@ -289,6 +291,8 @@ void quant_coarse_energy(const CELTMode *m, int start, int end, int effEnd,
       max_decay = MIN32(max_decay, .125f*nbAvailableBytes);
 #endif
    }
+   if (lfe)
+      max_decay=3;
    enc_start_state = *enc;
 
    ALLOC(oldEBands_intra, C*m->nbEBands, opus_val16);
@@ -298,7 +302,7 @@ void quant_coarse_energy(const CELTMode *m, int start, int end, int effEnd,
    if (two_pass || intra)
    {
       badness1 = quant_coarse_energy_impl(m, start, end, eBands, oldEBands_intra, budget,
-            tell, e_prob_model[LM][1], error_intra, enc, C, LM, 1, max_decay);
+            tell, e_prob_model[LM][1], error_intra, enc, C, LM, 1, max_decay, lfe);
    }
 
    if (!intra)
@@ -325,7 +329,7 @@ void quant_coarse_energy(const CELTMode *m, int start, int end, int effEnd,
       *enc = enc_start_state;
 
       badness2 = quant_coarse_energy_impl(m, start, end, eBands, oldEBands, budget,
-            tell, e_prob_model[LM][intra], error, enc, C, LM, 0, max_decay);
+            tell, e_prob_model[LM][intra], error, enc, C, LM, 0, max_decay, lfe);
 
       if (two_pass && (badness1 < badness2 || (badness1 == badness2 && ((opus_int32)ec_tell_frac(enc))+intra_bias > tell_intra)))
       {
diff --git a/celt/quant_bands.h b/celt/quant_bands.h
index b3187fad..0490bca4 100644
--- a/celt/quant_bands.h
+++ b/celt/quant_bands.h
@@ -51,7 +51,7 @@ void quant_coarse_energy(const CELTMode *m, int start, int end, int effEnd,
       const opus_val16 *eBands, opus_val16 *oldEBands, opus_uint32 budget,
       opus_val16 *error, ec_enc *enc, int C, int LM,
       int nbAvailableBytes, int force_intra, opus_val32 *delayedIntra,
-      int two_pass, int loss_rate);
+      int two_pass, int loss_rate, int lfe);
 
 void quant_fine_energy(const CELTMode *m, int start, int end, opus_val16 *oldEBands, opus_val16 *error, int *fine_quant, ec_enc *enc, int C);
 
diff --git a/src/opus_encoder.c b/src/opus_encoder.c
index 88bf5aff..235f5573 100644
--- a/src/opus_encoder.c
+++ b/src/opus_encoder.c
@@ -78,6 +78,7 @@ struct OpusEncoder {
     opus_int32   user_bitrate_bps;
     int          lsb_depth;
     int          encoder_buffer;
+    int          lfe;
 
 #define OPUS_ENCODER_RESET_START stream_channels
     int          stream_channels;
@@ -1234,6 +1235,11 @@ opus_int32 opus_encode_native(OpusEncoder *st, const opus_val16 *pcm, int frame_
     /* CELT mode doesn't support mediumband, use wideband instead */
     if (st->mode == MODE_CELT_ONLY && st->bandwidth == OPUS_BANDWIDTH_MEDIUMBAND)
         st->bandwidth = OPUS_BANDWIDTH_WIDEBAND;
+    if (st->lfe)
+    {
+       st->bandwidth = OPUS_BANDWIDTH_NARROWBAND;
+       st->mode = MODE_CELT_ONLY;
+    }
 
     /* Can't support higher than wideband for >20 ms frames */
     if (frame_size > st->Fs/50 && (st->mode == MODE_CELT_ONLY || st->bandwidth > OPUS_BANDWIDTH_WIDEBAND))
@@ -2203,6 +2209,13 @@ int opus_encoder_ctl(OpusEncoder *st, int request, ...)
             st->user_forced_mode = value;
         }
         break;
+        case OPUS_SET_LFE_REQUEST:
+        {
+            opus_int32 value = va_arg(ap, opus_int32);
+            st->lfe = value;
+            celt_encoder_ctl(celt_enc, OPUS_SET_LFE(value));
+        }
+        break;
 
         case CELT_GET_MODE_REQUEST:
         {
diff --git a/src/opus_multistream_encoder.c b/src/opus_multistream_encoder.c
index b3e0ccc4..9dcbb71c 100644
--- a/src/opus_multistream_encoder.c
+++ b/src/opus_multistream_encoder.c
@@ -329,7 +329,7 @@ static void surround_rate_allocation(
       int total = ((st->layout.nb_streams-st->layout.nb_coupled_streams-(st->lfe_stream!=-1))<<8) /* mono */
             + coupled_ratio*st->layout.nb_coupled_streams /* stereo */
             + (st->lfe_stream!=-1)*lfe_ratio;
-      channel_rate = 256*st->bitrate_bps/total;
+      channel_rate = 256*(st->bitrate_bps-2000)/total;
    }
 #ifndef FIXED_POINT
    if (st->variable_duration==OPUS_FRAMESIZE_VARIABLE && frame_size != Fs/50)
@@ -347,7 +347,7 @@ static void surround_rate_allocation(
       else if (i!=st->lfe_stream)
          rate[i] = channel_rate;
       else
-         rate[i] = channel_rate*lfe_ratio>>8;
+         rate[i] = 2000+(channel_rate*lfe_ratio>>8);
    }