From 824c752c51757753d35836e2c623e0161497a692 Mon Sep 17 00:00:00 2001
From: David Westgate <davidjwestgate@gmail.com>
Date: Sun, 10 Nov 2024 18:47:20 -0800
Subject: [PATCH] adaptive tone control script work

---
 code/adaptive-tone-control/main.py | 70 ++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 13 deletions(-)

diff --git a/code/adaptive-tone-control/main.py b/code/adaptive-tone-control/main.py
index 8751be6..388a857 100644
--- a/code/adaptive-tone-control/main.py
+++ b/code/adaptive-tone-control/main.py
@@ -1,20 +1,64 @@
-import string
 from scipy.io import wavfile as wav
-import scipy.fft as fft
 import numpy as np
-import sounddevice as sd
-import math 
-import scipy.signal as signal
+import sys
 
 print("Portfolio Object 2: Adaptive Tone Control")
-sample_rate, data = wav.read("song.wav")
 
-result = fft.fft(data)
+# Contants
+bands = {"low": (0, 300), "mid": (300, 2000), "high": (2000, 20000)}
 
-bands = {
-    "low":  ( 0,   300  ),
-    "mid":  ( 300, 2000 ),
-    "high": ( 2000,20000)
-}
+file = "test.wav"
+if len(sys.argv) > 1:
+    file = sys.argv[1]
+print("Input file: ", file)
+# Init
+sample_rate, wav_signal = wav.read(file)
+# window_size = 1024
+window_size = len(wav_signal)
+fft_values = np.fft.fft(wav_signal[:window_size])
+fft_freqs = np.fft.fftfreq(window_size, 1 / sample_rate)
 
-print('WIP')
\ No newline at end of file
+
+def band_energy(band, fft_values):
+    idx_band = np.where((fft_freqs >= band[0]) & (fft_freqs <= band[1]))[0]
+    return np.sum(np.abs(fft_values[idx_band]) ** 2)
+
+
+energy_low = band_energy(bands["low"], fft_values)
+energy_mid = band_energy(bands["mid"], fft_values)
+energy_high = band_energy(bands["high"], fft_values)
+avg_energy = (energy_low + energy_mid + energy_high) / 3
+
+
+print(f"low {energy_low:.2e}")
+print(f"mid {energy_mid:.2e}")
+print(f"high {energy_high:.2e}")
+print(f"avg {avg_energy:.2e}")
+
+
+def adjust(target_energy, current_energy, fft_values, band):
+    idx_band = np.where((fft_freqs >= band[0]) & (fft_freqs <= band[1]))[0]
+    gain = np.sqrt(target_energy / (current_energy + 1e-6))
+    adjusted_fft_values = np.copy(fft_values)
+    adjusted_fft_values[idx_band] *= gain
+    return adjusted_fft_values
+
+
+adjusted_low = adjust(avg_energy, energy_low, fft_values, bands["low"])
+adjusted_mid = adjust(avg_energy, energy_mid, fft_values, bands["mid"])
+adjusted_high = adjust(avg_energy, energy_high, fft_values, bands["high"])
+
+print(f'adj low {band_energy( bands["low"],adjusted_low):.2e}')
+print(f'adj mid {band_energy(bands["mid"],adjusted_mid):.2e}')
+print(f'adj high {band_energy( bands["high"],adjusted_high):.2e}')
+
+adjusted_all = adjusted_low + adjusted_mid + adjusted_high
+full_spectrum = np.concatenate([adjusted_all])
+reconstructed_signal = np.fft.ifft(adjusted_all).real
+reconstructed_signal = np.int16(
+    reconstructed_signal / np.max(np.abs(reconstructed_signal)) * np.iinfo(np.int16).max
+)
+
+output_file = "adj-" + file
+wav.write(output_file, sample_rate, reconstructed_signal)
+print(f"Adjusted audio written to {output_file}")