Attempt to use adaptive ATH for ATRAC1new_psy

5 files changed, 50 insertions, 8 deletions

diff --git a/src/atrac/atrac1_bitalloc.cpp b/src/atrac/atrac1_bitalloc.cpp
index a2e63fa..21584a9 100644
--- a/src/atrac/atrac1_bitalloc.cpp
+++ b/src/atrac/atrac1_bitalloc.cpp
@@ -127,12 +127,15 @@ vector<uint32_t> TAtrac1SimpleBitAlloc::CalcBitsAllocation(const std::vector<TSc
                                                            const uint32_t bfuNum,
                                                            const TFloat spread,
                                                            const TFloat shift,
-                                                           const TBlockSize& blockSize) {
+                                                           const TBlockSize& blockSize,
+                                                           const float loudness) {
     vector<uint32_t> bitsPerEachBlock(bfuNum);
     for (size_t i = 0; i < bitsPerEachBlock.size(); ++i) {
         bool shortBlock = blockSize.LogCount[BfuToBand(i)];
         const uint32_t fix = shortBlock ? FixedBitAllocTableShort[i] : FixedBitAllocTableLong[i];
-        if (!shortBlock && scaledBlocks[i].MaxEnergy < ATHLong[i]) {
+        float ath = ATHLong[i] * loudness;
+        //std::cerr << "block: " << i << " Loudness: " << loudness << " " << 10 * log10(scaledBlocks[i].MaxEnergy / ath) << std::endl;
+        if (!shortBlock && scaledBlocks[i].MaxEnergy < ath) {
             bitsPerEachBlock[i] = 0;
         } else {
             int tmp = spread * ( (TFloat)scaledBlocks[i].ScaleFactorIndex/3.2) + (1.0 - spread) * fix - shift;
@@ -183,7 +186,7 @@ uint32_t TAtrac1SimpleBitAlloc::CheckBfuUsage(bool* changed,
     return curBfuId;
 }
 
-uint32_t TAtrac1SimpleBitAlloc::Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize) {
+uint32_t TAtrac1SimpleBitAlloc::Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize, float loudness) {
     uint32_t bfuIdx = BfuIdxConst ? BfuIdxConst - 1 : 7;
     bool autoBfu = !BfuIdxConst;
     TFloat spread = AnalizeScaleFactorSpread(scaledBlocks);
@@ -204,7 +207,7 @@ uint32_t TAtrac1SimpleBitAlloc::Write(const std::vector<TScaledBlock>& scaledBlo
         bool bfuNumChanged = false;
         for (;;) {
             const vector<uint32_t>& tmpAlloc = CalcBitsAllocation(scaledBlocks, BfuAmountTab[bfuIdx],
-                                                                  spread, shift, blockSize);
+                                                                  spread, shift, blockSize, loudness);
             uint32_t bitsUsed = 0;
             for (size_t i = 0; i < tmpAlloc.size(); i++) {
                 bitsUsed += SpecsPerBlock[i] * tmpAlloc[i];
diff --git a/src/atrac/atrac1_bitalloc.h b/src/atrac/atrac1_bitalloc.h
index b1f6b54..4c541b8 100644
--- a/src/atrac/atrac1_bitalloc.h
+++ b/src/atrac/atrac1_bitalloc.h
@@ -33,7 +33,7 @@ class IAtrac1BitAlloc {
 public:
     IAtrac1BitAlloc() {};
     virtual ~IAtrac1BitAlloc() {};
-    virtual uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize) = 0;
+    virtual uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize, float loudness) = 0;
 };
 
 class TBitsBooster : public virtual TAtrac1Data {
@@ -56,7 +56,8 @@ public:
 
 class TAtrac1SimpleBitAlloc : public TAtrac1BitStreamWriter, public TBitsBooster, public virtual IAtrac1BitAlloc {
     std::vector<uint32_t> CalcBitsAllocation(const std::vector<TScaledBlock>& scaledBlocks, const uint32_t bfuNum,
-                                             const TFloat spread, const TFloat shift, const TBlockSize& blockSize);
+                                             const TFloat spread, const TFloat shift, const TBlockSize& blockSize,
+                                             const float loudness);
     const uint32_t BfuIdxConst;
     const bool FastBfuNumSearch;
     static std::vector<TFloat> ATHLong;
@@ -66,7 +67,7 @@ class TAtrac1SimpleBitAlloc : public TAtrac1BitStreamWriter, public TBitsBooster
 public:
     TAtrac1SimpleBitAlloc(ICompressedOutput* container, uint32_t bfuIdxConst, bool fastBfuNumSearch);
     ~TAtrac1SimpleBitAlloc() {};
-    uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize) override;
+    uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize, float loudness) override;
 };
 
 } //namespace NAtrac1
diff --git a/src/atrac/atrac_psy_common.cpp b/src/atrac/atrac_psy_common.cpp
index 089bf47..18f855c 100644
--- a/src/atrac/atrac_psy_common.cpp
+++ b/src/atrac/atrac_psy_common.cpp
@@ -135,4 +135,40 @@ vector<float> CalcATH(int len, int sampleRate)
     return res;
 }
 
+float TrackLoudness(float prevLoud, const TFloat* e0, const TFloat* e1, const float* weight, size_t sz)
+{
+    float s = 0;
+    if (e1 != nullptr) {
+        for (size_t i = 0; i < sz; i++) {
+            s += (e0[i] + e1[i]) * weight[i];
+        }
+
+        s *= 0.5;
+
+    } else {
+        for (size_t i = 0; i < sz; i++) {
+            s += e0[i] * weight[i];
+        }
+    }
+
+    return 0.98 * prevLoud + 0.02 * s;
+}
+
+vector<float> CreateLoudnessCurve(size_t sz)
+{
+    std::vector<float> res;
+    res.resize(sz);
+
+    for (size_t i = 0; i < sz; i++) {
+        float f = (float)(i + 3) * 0.5 * 44100 / (float)sz;
+        float t = std::log10(f) - 3.5;
+        t = -10 * t * t + 3 - f / 3000;
+        t = std::pow(10, (0.1 * t));
+        //std::cerr << i << "  => " << f << "  "  << t <<std::endl;
+        res[i] = t;
+    }
+
+    return res;
+}
+
 } // namespace NAtracDEnc
diff --git a/src/atrac/atrac_psy_common.h b/src/atrac/atrac_psy_common.h
index 1a826bb..ffd1f03 100644
--- a/src/atrac/atrac_psy_common.h
+++ b/src/atrac/atrac_psy_common.h
@@ -23,5 +23,7 @@ namespace NAtracDEnc {
 
 TFloat AnalizeScaleFactorSpread(const std::vector<TScaledBlock>& scaledBlocks);
 std::vector<float> CalcATH(int len, int sampleRate);
+float TrackLoudness(float prevLoud, const TFloat* e0, const TFloat* e1, const float* weight, size_t sz);
+std::vector<float> CreateLoudnessCurve(size_t sz);
 
 } //namespace NAtracDEnc
diff --git a/src/atrac/atrac_scale.h b/src/atrac/atrac_scale.h
index b590662..ec83059 100644
--- a/src/atrac/atrac_scale.h
+++ b/src/atrac/atrac_scale.h
@@ -77,7 +77,7 @@ public:
     TBlockSize()
         : LogCount({{0, 0, 0, 0}})
     {}
-    const std::array<int, 4> LogCount;
+    std::array<int, 4> LogCount;
 };
 
 } //namespace NAtracDEnc