feat: add power of 2 length, spectrogram example

Author: flavioschneider

IMAGE.png

@@ -1,6 +1,7 @@
 
 # CQT - PyTorch
 
+
 An invertible and differentiable implementation of the Constant-Q Transform (CQT) using Non-stationary Gabor Transform (NSGT), in PyTorch.
 
 ```bash
@@ -27,7 +28,11 @@ z = transform.encode(x) # [1, 2, 455, 2796] = [batch_size, channels, frequencies
 y = transform.decode(z) # [1, 1, 262144]
 ```
 
+### Example CQT spectrogram (z)
+<img src="./IMAGE.png"></img>
+
 ## TODO
+* [x] Power of 2 length (with `power_of_2_length` constructor arg).
 * [ ] Understand why/if inverse window is necessary.
 * [ ] Allow variable audio lengths by chunking.
 

README.md

@@ -1,10 +1,14 @@
-from math import floor
+from math import ceil, floor
 
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
 
 
+def next_power_of_2(x: Tensor) -> int:
+    return 2 ** ceil(x.item()).bit_length()
+
+
 def get_center_frequencies(
     num_octaves: int, num_bins_per_octave: int, sample_rate: int  # C  # B  # Xi_s
 ) -> Tensor:  # Xi_k for k in [1, 2*K+1]
@@ -51,25 +55,27 @@ def get_bandwidths(
     return bandwidths_all
 
 
-def get_windows_range_indices(lengths: Tensor, positions: Tensor) -> Tensor:
+def get_windows_range_indices(
+    lengths: Tensor, positions: Tensor, power_of_2_length: bool
+) -> Tensor:
     """Compute windowing tensor of indices"""
     num_bins = lengths.shape[0] // 2
-    max_length = lengths.max()
+    max_length = next_power_of_2(lengths.max()) if power_of_2_length else lengths.max()
     ranges = []
     for i in range(num_bins):
         start = positions[i] - max_length
         ranges += [torch.arange(start=start, end=start + max_length)]  # type: ignore
     return torch.stack(ranges, dim=0).long()
 
 
-def get_windows(lengths: Tensor) -> Tensor:
+def get_windows(lengths: Tensor, power_of_2_length: bool) -> Tensor:
     """Compute tensor of stacked (centered) windows"""
     num_bins = lengths.shape[0] // 2
-    max_length = lengths.max()
+    max_length = next_power_of_2(lengths.max()) if power_of_2_length else lengths.max()
     windows = []
     for length in lengths[:num_bins]:
         # Pad windows left and right to center them
-        pad_left = floor(max_length / 2 - length / 2)
+        pad_left = floor(max_length / 2 - length / 2)  # type: ignore
         pad_right = int(max_length - length - pad_left)
         windows += [F.pad(torch.hann_window(int(length)), pad=(pad_left, pad_right))]
     return torch.stack(windows, dim=0)
@@ -87,6 +93,7 @@ def __init__(
         num_bins_per_octave: int,
         sample_rate: int,
         block_length: int,
+        power_of_2_length: bool = False,
     ):
         super().__init__()
         self.block_length = block_length
@@ -111,10 +118,14 @@ def __init__(
             get_windows_range_indices(
                 lengths=window_lengths,
                 positions=torch.round(frequencies * block_length / sample_rate),
+                power_of_2_length=power_of_2_length,
             ),
         )
 
-        self.register_buffer("windows", get_windows(lengths=window_lengths))
+        self.register_buffer(
+            "windows",
+            get_windows(lengths=window_lengths, power_of_2_length=power_of_2_length),
+        )
 
         self.register_buffer(
             "windows_inverse",

cqt_pytorch/cqt.py

@@ -3,7 +3,7 @@
 setup(
     name="cqt-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.2",
+    version="0.0.3",
     license="MIT",
     description="CQT Pytorch",
     long_description_content_type="text/markdown",