cleanup python loudness script

script/louder_wav.py

@@ -16,18 +16,38 @@ SHRT_MAX = 32767
 # new_frames = struct.pack('h'*len(frames), *frames)
 
 def read_wav(filename):
-    """Read a wav file, return a tuple of (wav_params, frames)."""
+    """
+    Read a wav file, return a tuple of (wav_params, frames_bytestring).
+    """
     params = None
     frames = None
     with wave.open(filename, 'rb') as w:
         params = w.getparams()
         numframes = params[3]
-        frames = w.readframes(numframes)
-    return (params, frames)
+        frames_bytestring = w.readframes(numframes)
+    return (params, frames_bytestring)
+
+def write_wav(output_filename, wav_params, frames_bytestring):
+    """Write a wav file."""
+    with wave.open(output_filename, 'wb') as w:
+        w.setparams(wav_params)
+        w.writeframes(frames_bytestring)
 
 def generate_noise():
     pass
 
+def get_struct_format(wav_params, num_samples):
+    sample_width = wav_params[1]
+    if sample_width not in [1, 2]:
+        raise ValueError("Only supports 8 and 16 bit audio.")
+    if sample_width == 1:
+        # read unsigned chars
+        return '{}B'.format(num_samples)
+    else:
+        # read signed 2 byte shorts
+        return '{}h'.format(num_samples)
+
+
 def get_pcm_channels(wav_params, frames):
     """
     Decompose a wave bytestring (with its associated metadata) into its
@@ -47,14 +67,9 @@ def get_pcm_channels(wav_params, frames):
 
     total_samples = num_frames * num_channels
 
-    if sample_width == 1:
-        # read unsigned chars
-        fmt = '{}B'.format(total_samples)
-    else:
-        # read signed 2 byte shorts
-        fmt = '{}h'.format(total_samples)
-
+    fmt = get_struct_format(wav_params, total_samples)
     integer_data = struct.unpack(fmt, frames)
+
     channels = [[] for i in range(num_channels)]
 
     for i, value in enumerate(integer_data):
@@ -63,7 +78,8 @@ def get_pcm_channels(wav_params, frames):
 
     return channels
 
-def channels_to_bytestream(wav_params, pcm_data):
+def channels_to_bytestring(wav_params, pcm_data):
+    sample_width = wav_params[1]
     num_channels = wav_params[0]
     num_frames = wav_params[3]
 
@@ -73,24 +89,32 @@ def channels_to_bytestream(wav_params, pcm_data):
     integer_data = []
 
     for j in range(num_frames):
-        integer = 0
         for i in range(num_channels):
-            integer = pcm_data[i][j]
-
-            # Change volume by changing the signal amplitude!
-            integer *= 10
-
-            # Don't overflow.
-            short = saturate_signed_short(integer)
-
-            integer_data.append(short)
+            integer_data.append(pcm_data[i][j])
 
     print("Length of pcm_data: {}".format(len(pcm_data)))
     print("Length of a pcm channel: {}".format(len(pcm_data[0])))
     print("Length of packed integer data: {}".format(len(integer_data)))
 
-    return struct.pack('{}h'.format(num_frames*2), *integer_data)
+    total_samples = num_frames * num_channels
 
+    fmt = get_struct_format(wav_params, total_samples)
+    return struct.pack(fmt, *integer_data)
+
+def increase_volume(pcm_data, multiplier=2):
+    """
+    Increase the volume of the PCM data.
+    """
+    num_channels = len(pcm_data)
+    num_frames = len(pcm_data[0])
+
+    for j in range(num_frames):
+        for i in range(num_channels):
+            x = pcm_data[i][j]
+            # Change volume by changing the signal amplitude!
+            x *= multiplier
+            # Don't overflow.
+            pcm_data[i][j] = saturate_signed_short(x)
 
 def saturate_signed_short(integer):
     """Don't overflow the signed short width."""
@@ -98,12 +122,7 @@ def saturate_signed_short(integer):
         return SHRT_MAX
     elif integer < SHRT_MIN:
         return SHRT_MIN
-    return integer
-
-def write_wav(output_filename, wav_params, frames):
-    with wave.open(output_filename, 'wb') as w:
-        w.setparams(wav_params)
-        w.writeframes(frames)
+    return int(integer)
 
 def main():
     if len(sys.argv) < 3:
@@ -111,12 +130,19 @@ def main():
     infile = sys.argv[1]
     outfile = sys.argv[2]
 
+    multiplier = 2
+    if len(sys.argv) > 3:
+        multiplier = int(sys.argv[3])
+
     params, frames = read_wav(infile)
     print(params)
 
     channels = get_pcm_channels(params, frames)
 
-    bytestring = channels_to_bytestream(params, channels)
+    print('Volume multiplier: {}'.format(multiplier))
+    increase_volume(channels, multiplier)
+
+    bytestring = channels_to_bytestring(params, channels)
 
     print("Length of original bytestring: {}".format(len(frames)))
     print("Length of new bytestring: {}".format(len(bytestring)))