Merge branch 'v0.1.5' of github.com:v0lta/PyTorch-Wavelet-Toolbox into v0.1.5

Author: v0lta

README.rst

@@ -45,8 +45,8 @@ Welcome to the PyTorch wavelet toolbox. This package implements:
 - single and two-dimensional wavelet packet forward and backward transforms are available via the ``WaveletPacket`` and ``WaveletPacket2D`` objects,
 - finally, this package provides adaptive wavelet support (experimental).
 
-This toolbox supports pywt-wavelets. Complete documentation is available:
-https://pytorch-wavelet-toolbox.readthedocs.io/
+This toolbox extends `PyWavelets <https://pywavelets.readthedocs.io/en/latest/>`_ . We additionally provide GPU and gradient support via a PyTorch backend.
+Complete documentation is available at: https://pytorch-wavelet-toolbox.readthedocs.io/
 
 
 **Installation**
@@ -149,8 +149,8 @@ Reconsidering the 1d case, try:
 
 
 The process for the 2d transforms ``MatrixWavedec2``, ``MatrixWaverec2`` works similarly.
-By default, a non-separable transformation is used.
-To use a separable transformation, pass ``separable=True`` to ``MatrixWavedec2`` and ``MatrixWaverec2``.
+By default, a separable transformation is used.
+To use a non-separable transformation, pass ``separable=False`` to ``MatrixWavedec2`` and ``MatrixWaverec2``.
 Separable transformations use a 1d transformation along both axes, which might be faster since fewer matrix entries
 have to be orthogonalized.
 

examples/deepfake_analysis/README.md

@@ -25,7 +25,7 @@ Recreating these experiments requires roughly 8GB of free disc space and 10GB of
 To reproduce these plots:
 1. Download [ffhq_style_gan.zip](https://drive.google.com/uc?id=1MOHKuEVqURfCKAN9dwp1o2tuR19OTQCF&export=download) and
 2. Extract the image pairs here.
-3. Check the file structure. In the `ffhq_style_gan` the folder structure should be:
+3. Check the file structure. In `ffhq_style_gan` the folder structure should be:
 ```
 source_data
  ├── A_ffhq

src/ptwt/__init__.py

@@ -8,9 +8,4 @@
 from .matmul_transform_2 import MatrixWavedec2, MatrixWaverec2
 from .matmul_transform_3 import MatrixWavedec3, MatrixWaverec3
 from .packets import WaveletPacket, WaveletPacket2D
-from .separable_conv_transform import (
-    fswavedec2,
-    fswavedec3,
-    fswaverec2,
-    fswaverec3,
-)
+from .separable_conv_transform import fswavedec2, fswavedec3, fswaverec2, fswaverec3

src/ptwt/conv_transform.py

@@ -186,6 +186,20 @@ def _flatten_2d_coeff_lst(
     return flat_coeff_lst
 
 
+def _adjust_padding_at_reconstruction(
+    res_ll_size: int, coeff_size: int, pad_end: int, pad_start: int
+) -> Tuple[int, int]:
+    pred_size = res_ll_size - (pad_start + pad_end)
+    next_size = coeff_size
+    if next_size == pred_size:
+        pass
+    elif next_size == pred_size - 1:
+        pad_end += 1
+    else:
+        raise AssertionError("padding error, please open an issue on github")
+    return pad_end, pad_start
+
+
 def wavedec(
     data: torch.Tensor,
     wavelet: Union[Wavelet, str],
@@ -318,14 +332,9 @@ def waverec(coeffs: List[torch.Tensor], wavelet: Union[Wavelet, str]) -> torch.T
         padl = (2 * filt_len - 3) // 2
         padr = (2 * filt_len - 3) // 2
         if c_pos < len(coeffs) - 2:
-            pred_len = res_lo.shape[-1] - (padl + padr)
-            next_len = coeffs[c_pos + 2].shape[-1]
-            if next_len != pred_len:
-                padr += 1
-                pred_len = res_lo.shape[-1] - (padl + padr)
-                assert (
-                    next_len == pred_len
-                ), "padding error, please open an issue on github "
+            padr, padl = _adjust_padding_at_reconstruction(
+                res_lo.shape[-1], coeffs[c_pos + 2].shape[-1], padr, padl
+            )
         if padl > 0:
             res_lo = res_lo[..., padl:]
         if padr > 0:

src/ptwt/conv_transform_2.py

@@ -11,7 +11,12 @@
 import torch
 
 from ._util import Wavelet, _as_wavelet, _get_len, _is_dtype_supported, _outer
-from .conv_transform import _get_pad, _translate_boundary_strings, get_filter_tensors
+from .conv_transform import (
+    _adjust_padding_at_reconstruction,
+    _get_pad,
+    _translate_boundary_strings,
+    get_filter_tensors,
+)
 
 
 def construct_2d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
@@ -244,22 +249,13 @@ def waverec2(
         padt = (2 * filt_len - 3) // 2
         padb = (2 * filt_len - 3) // 2
         if c_pos < len(coeffs) - 2:
-            pred_len = res_ll.shape[-1] - (padl + padr)
-            next_len = coeffs[c_pos + 2][0].shape[-1]
-            pred_len2 = res_ll.shape[-2] - (padt + padb)
-            next_len2 = coeffs[c_pos + 2][0].shape[-2]
-            if next_len != pred_len:
-                padr += 1
-                pred_len = res_ll.shape[-1] - (padl + padr)
-                assert (
-                    next_len == pred_len
-                ), "padding error, please open an issue on github "
-            if next_len2 != pred_len2:
-                padb += 1
-                pred_len2 = res_ll.shape[-2] - (padt + padb)
-                assert (
-                    next_len2 == pred_len2
-                ), "padding error, please open an issue on github "
+            padr, padl = _adjust_padding_at_reconstruction(
+                res_ll.shape[-1], coeffs[c_pos + 2][0].shape[-1], padr, padl
+            )
+            padb, padt = _adjust_padding_at_reconstruction(
+                res_ll.shape[-2], coeffs[c_pos + 2][0].shape[-2], padb, padt
+            )
+
         if padt > 0:
             res_ll = res_ll[..., padt:, :]
         if padb > 0:

src/ptwt/conv_transform_3.py

@@ -3,13 +3,18 @@
 The functions here are based on torch.nn.functional.conv3d and it's transpose.
 """
 
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Sequence, Union, cast
 
 import pywt
 import torch
 
 from ._util import Wavelet, _as_wavelet, _get_len, _is_dtype_supported, _outer
-from .conv_transform import _get_pad, _translate_boundary_strings, get_filter_tensors
+from .conv_transform import (
+    _adjust_padding_at_reconstruction,
+    _get_pad,
+    _translate_boundary_strings,
+    get_filter_tensors,
+)
 
 
 def _construct_3d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
@@ -153,7 +158,7 @@ def wavedec3(
 
 
 def waverec3(
-    coeffs: List[Union[torch.Tensor, Dict[str, torch.Tensor]]],
+    coeffs: Sequence[Union[torch.Tensor, Dict[str, torch.Tensor]]],
     wavelet: Union[Wavelet, str],
 ) -> torch.Tensor:
     """Reconstruct a signal from wavelet coefficients.
@@ -199,7 +204,8 @@ def waverec3(
     filt_len = rec_lo.shape[-1]
     rec_filt = _construct_3d_filt(lo=rec_lo, hi=rec_hi)
 
-    for c_pos, coeff_dict in enumerate(coeffs[1:]):
+    coeff_dicts = cast(Sequence[Dict[str, torch.Tensor]], coeffs[1:])
+    for c_pos, coeff_dict in enumerate(coeff_dicts):
         if not isinstance(coeff_dict, dict) or len(coeff_dict) != 7:
             raise ValueError(
                 f"Unexpected detail coefficient type: {type(coeff_dict)}. Detail "
@@ -238,31 +244,16 @@ def waverec3(
         padr = (2 * filt_len - 3) // 2
         padt = (2 * filt_len - 3) // 2
         padb = (2 * filt_len - 3) // 2
-        if c_pos < len(coeffs) - 2:
-            pred_len = res_lll.shape[-1] - (padl + padr)
-            next_len = coeffs[c_pos + 2]["aad"].shape[-1]  # type: ignore
-            pred_len2 = res_lll.shape[-2] - (padt + padb)
-            next_len2 = coeffs[c_pos + 2]["aad"].shape[-2]  # type: ignore
-            pred_len3 = res_lll.shape[-3] - (padfr + padba)
-            next_len3 = coeffs[c_pos + 2]["aad"].shape[-3]  # type: ignore
-            if next_len != pred_len:
-                padr += 1
-                pred_len = res_lll.shape[-1] - (padl + padr)
-                assert (
-                    next_len == pred_len
-                ), "padding error, please open an issue on github "
-            if next_len2 != pred_len2:
-                padb += 1
-                pred_len2 = res_lll.shape[-2] - (padt + padb)
-                assert (
-                    next_len2 == pred_len2
-                ), "padding error, please open an issue on github "
-            if next_len3 != pred_len3:
-                padba += 1
-                pred_len3 = res_lll.shape[-3] - (padba + padfr)
-                assert (
-                    next_len3 == pred_len3
-                ), "padding error, please open an issue on github "
+        if c_pos + 1 < len(coeff_dicts):
+            padr, padl = _adjust_padding_at_reconstruction(
+                res_lll.shape[-1], coeff_dicts[c_pos + 1]["aad"].shape[-1], padr, padl
+            )
+            padb, padt = _adjust_padding_at_reconstruction(
+                res_lll.shape[-2], coeff_dicts[c_pos + 1]["aad"].shape[-2], padb, padt
+            )
+            padba, padfr = _adjust_padding_at_reconstruction(
+                res_lll.shape[-3], coeff_dicts[c_pos + 1]["aad"].shape[-3], padba, padfr
+            )
         if padt > 0:
             res_lll = res_lll[..., padt:, :]
         if padb > 0:

src/ptwt/packets.py

@@ -4,7 +4,7 @@
 import collections
 from functools import partial
 from itertools import product
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import pywt
 import torch
@@ -14,6 +14,7 @@
 from .conv_transform_2 import wavedec2, waverec2
 from .matmul_transform import MatrixWavedec, MatrixWaverec
 from .matmul_transform_2 import MatrixWavedec2, MatrixWaverec2
+from .separable_conv_transform import fswavedec2, fswaverec2
 
 if TYPE_CHECKING:
     BaseDict = collections.UserDict[str, torch.Tensor]
@@ -129,6 +130,12 @@ def reconstruct(self) -> "WaveletPacket":
                 data_a = self[node + "a"]
                 data_b = self[node + "d"]
                 rec = self._get_waverec(data_a.shape[-1])([data_a, data_b])
+                if level > 0:
+                    if rec.shape[-1] != self[node].shape[-1]:
+                        assert (
+                            rec.shape[-1] == self[node].shape[-1] + 1
+                        ), "padding error, please open an issue on github"
+                        rec = rec[..., :-1]
                 self[node] = rec
         return self
 
@@ -249,9 +256,8 @@ def __init__(
                 to use in the sparse matrix backend. Only used if `mode`
                 equals 'boundary'. Choose from 'qr' or 'gramschmidt'.
                 Defaults to 'qr'.
-            separable (bool): If true and the sparse matrix backend is selected,
-                a separable transform is performed, i.e. each image axis is
-                transformed separately. Defaults to False.
+            separable (bool): If true, a separable transform is performed,
+                i.e. each image axis is transformed separately. Defaults to False.
             maxlevel (int, optional): Value is passed on to `transform`.
                 The highest decomposition level to compute. If None, the maximum level
                 is determined from the input data shape. Defaults to None.
@@ -288,6 +294,10 @@ def transform(
             maxlevel = pywt.dwt_max_level(min(data.shape[-2:]), self.wavelet.dec_len)
         self.maxlevel = maxlevel
 
+        if data.dim() == 2:
+            # add batch dim to unbatched input
+            data = data.unsqueeze(0)
+
         self._recursive_dwt2d(data, level=0, path="")
         return self
 
@@ -306,24 +316,25 @@ def reconstruct(self) -> "WaveletPacket2D":
 
         for level in reversed(range(self.maxlevel)):
             for node in self.get_natural_order(level):
-                if self.mode == "boundary":
-                    data_a = self[node + "a"]
-                    data_h = self[node + "h"]
-                    data_v = self[node + "v"]
-                    data_d = self[node + "d"]
-                    rec = self._get_waverec(data_a.shape[-2:])(
-                        (data_a, (data_h, data_v, data_d))
-                    )
-                    self[node] = rec
-                else:
-                    data_a = self[node + "a"]
-                    data_h = self[node + "h"]
-                    data_v = self[node + "v"]
-                    data_d = self[node + "d"]
-                    rec = self._get_waverec(data_a.shape[-2:])(
-                        (data_a, (data_h, data_v, data_d))
-                    )
-                    self[node] = rec.squeeze(1)
+                data_a = self[node + "a"]
+                data_h = self[node + "h"]
+                data_v = self[node + "v"]
+                data_d = self[node + "d"]
+                rec = self._get_waverec(data_a.shape[-2:])(
+                    [data_a, (data_h, data_v, data_d)]
+                )
+                if level > 0:
+                    if rec.shape[-1] != self[node].shape[-1]:
+                        assert (
+                            rec.shape[-1] == self[node].shape[-1] + 1
+                        ), "padding error, please open an issue on github"
+                        rec = rec[..., :-1]
+                    if rec.shape[-2] != self[node].shape[-2]:
+                        assert (
+                            rec.shape[-2] == self[node].shape[-2] + 1
+                        ), "padding error, please open an issue on github"
+                        rec = rec[..., :-1, :]
+                self[node] = rec
         return self
 
     def get_natural_order(self, level: int) -> List[str]:
@@ -354,13 +365,17 @@ def _get_wavedec(
                 )
             fun = self.matrix_wavedec2_dict[shape]
             return fun
+        elif self.separable:
+            return self._transform_fsdict_to_tuple_func(
+                partial(fswavedec2, wavelet=self.wavelet, level=1, mode=self.mode)
+            )
         else:
             return partial(wavedec2, wavelet=self.wavelet, level=1, mode=self.mode)
 
     def _get_waverec(
         self, shape: Tuple[int, ...]
     ) -> Callable[
-        [Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]],
+        [List[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]],
         torch.Tensor,
     ]:
         if self.mode == "boundary":
@@ -371,11 +386,54 @@ def _get_waverec(
                     boundary=self.boundary,
                     separable=self.separable,
                 )
-            fun = self.matrix_waverec2_dict[shape]
-            return fun  # type: ignore
+            return self.matrix_waverec2_dict[shape]
+        elif self.separable:
+            return self._transform_tuple_to_fsdict_func(
+                partial(fswaverec2, wavelet=self.wavelet)
+            )
         else:
             return partial(waverec2, wavelet=self.wavelet)
 
+    def _transform_fsdict_to_tuple_func(
+        self,
+        fs_dict_func: Callable[
+            [torch.Tensor], List[Union[torch.Tensor, Dict[str, torch.Tensor]]]
+        ],
+    ) -> Callable[
+        [torch.Tensor],
+        List[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]],
+    ]:
+        def _tuple_func(
+            data: torch.Tensor,
+        ) -> List[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]:
+            a_coeff, fsdict = fs_dict_func(data)
+            fsdict = cast(Dict[str, torch.Tensor], fsdict)
+            return [
+                cast(torch.Tensor, a_coeff),
+                (fsdict["ad"], fsdict["da"], fsdict["dd"]),
+            ]
+
+        return _tuple_func
+
+    def _transform_tuple_to_fsdict_func(
+        self,
+        fsdict_func: Callable[
+            [List[Union[torch.Tensor, Dict[str, torch.Tensor]]]], torch.Tensor
+        ],
+    ) -> Callable[
+        [List[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]],
+        torch.Tensor,
+    ]:
+        def _fsdict_func(
+            coeffs: List[
+                Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
+            ]
+        ) -> torch.Tensor:
+            a, (h, v, d) = coeffs
+            return fsdict_func([cast(torch.Tensor, a), {"ad": h, "da": v, "dd": d}])
+
+        return _fsdict_func
+
     def _recursive_dwt2d(self, data: torch.Tensor, level: int, path: str) -> None:
         if not self.maxlevel:
             raise AssertionError

tests/test_separable_conv_fwt.py

@@ -9,12 +9,12 @@
 from src.ptwt.matmul_transform_2 import MatrixWavedec2
 from src.ptwt.matmul_transform_3 import MatrixWavedec3
 from src.ptwt.separable_conv_transform import (
+    _fswavedec,
+    _fswaverec,
     _separable_conv_wavedecn,
     _separable_conv_waverecn,
-    _fswavedec,
     fswavedec2,
     fswavedec3,
-    _fswaverec,
     fswaverec2,
     fswaverec3,
 )