Support spectrograms

ultratrace2/model/files/__init__.py

@@ -68,3 +68,12 @@
     )
 except ImportError as e:
     logger.warning(e)
+
+try:
+    from .loaders import ParselmouthLoader
+
+    __register(
+        [".pmpkl"], ["application/octet-stream"], ParselmouthLoader,
+    )
+except ImportError as e:
+    logger.warning(e)

ultratrace2/model/files/bundle.py

@@ -7,9 +7,11 @@
     AlignmentFileLoader,
     ImageSetFileLoader,
     SoundFileLoader,
+    SpectrogramFileLoader,
     FileLoaderBase,
     FileLoadError,
 )
+from .loaders.parselmouth_spectrogram import ParselmouthLoader
 from .registry import get_loader_for
 
 
@@ -23,35 +25,79 @@ def __init__(
         alignment_file: Optional[AlignmentFileLoader] = None,
         image_set_file: Optional[ImageSetFileLoader] = None,
         sound_file: Optional[SoundFileLoader] = None,
+        spectrogram_file: Optional[SpectrogramFileLoader] = None,
     ):
         self.name = name
         self.alignment_file = alignment_file
         self.image_set_file = image_set_file
         self.sound_file = sound_file
+        self.spectrogram_file = spectrogram_file
 
     def has_impl(self) -> bool:
         return any(
             f is not None
-            for f in [self.alignment_file, self.image_set_file, self.sound_file]
+            for f in [
+                self.alignment_file,
+                self.image_set_file,
+                self.sound_file,
+                self.spectrogram_file,
+            ]
         )
 
+    def get_alignment_file(self) -> Optional[AlignmentFileLoader]:
+        return self.alignment_file
+
     def set_alignment_file(self, alignment_file: AlignmentFileLoader) -> None:
         if self.alignment_file is not None:
             logger.warning("Overwriting existing alignment file")
         self.alignment_file = alignment_file
 
+    def get_image_set_file(self) -> Optional[ImageSetFileLoader]:
+        return self.image_set_file
+
     def set_image_set_file(self, image_set_file: ImageSetFileLoader) -> None:
         if self.image_set_file is not None:
             logger.warning("Overwriting existing image-set file")
         self.image_set_file = image_set_file
 
+    def get_sound_file(self) -> Optional[SoundFileLoader]:
+        return self.sound_file
+
     def set_sound_file(self, sound_file: SoundFileLoader) -> None:
         if self.sound_file is not None:
             logger.warning("Overwriting existing sound file")
         self.sound_file = sound_file
 
+    def get_spectrogram_file(self) -> Optional[SpectrogramFileLoader]:
+        if self.spectrogram_file is None:
+            sound_file = self.get_sound_file()
+            if sound_file is not None:
+                try:
+                    logger.debug(
+                        "No spectrogram file found, trying to create from sound file"
+                    )
+                    self.set_spectrogram_file(
+                        ParselmouthLoader.from_sound_file(sound_file)
+                    )
+                except FileLoadError as e:
+                    logger.warning(e)
+
+        return self.spectrogram_file
+
+    def set_spectrogram_file(self, spectrogram_file: SpectrogramFileLoader) -> None:
+        if self.spectrogram_file is not None:
+            logger.warning("Overwriting existing spectrogram file")
+        self.spectrogram_file = spectrogram_file
+
     def __repr__(self):
-        return f'Bundle("{self.name}",{self.alignment_file},{self.image_set_file},{self.sound_file})'
+        return (
+            f'Bundle("{self.name}"'
+            f",{self.get_alignment_file()}"
+            f",{self.get_image_set_file()}"
+            f",{self.get_sound_file()}"
+            f",{self.get_spectrogram_file()}"
+            f")"
+        )
 
 
 class FileBundleList:
@@ -74,11 +120,13 @@ def __init__(self, bundles: Mapping[str, FileBundle]):
         self.has_alignment_impl: bool = False
         self.has_image_set_impl: bool = False
         self.has_sound_impl: bool = False
+        self.has_spectrogram_impl: bool = False
 
         for bundle in bundles.values():
-            self.has_alignment_impl |= bundle.alignment_file is not None
-            self.has_image_set_impl |= bundle.image_set_file is not None
-            self.has_sound_impl |= bundle.sound_file is not None
+            self.has_alignment_impl |= bundle.get_alignment_file() is not None
+            self.has_image_set_impl |= bundle.get_image_set_file() is not None
+            self.has_sound_impl |= bundle.get_sound_file() is not None
+            self.has_spectrogram_impl |= bundle.get_spectrogram_file() is not None
 
     @classmethod
     def build_from_dir(

ultratrace2/model/files/loaders/__init__.py

@@ -10,3 +10,6 @@
 from .mp3 import MP3Loader  # noqa: F401
 from .ogg import OggLoader  # noqa: F401
 from .wav import WAVLoader  # noqa: F401
+
+# spectrogram files
+from .parselmouth_spectrogram import ParselmouthLoader  # noqa: F401

ultratrace2/model/files/loaders/base.py

@@ -108,3 +108,22 @@ class SoundFileLoader(FileLoaderBase):
     @abstractmethod
     def __len__(self) -> int:
         """Length of file in ms"""
+
+
+class SpectrogramFileLoader(FileLoaderBase):
+    @classmethod
+    @abstractmethod
+    def from_sound_file(cls: Type[Self], sound_file: SoundFileLoader) -> Self:
+        ...
+
+    @abstractmethod
+    def get_image(
+        self,
+        start_time_ms: int,
+        stop_time_ms: int,
+        window_length: float,
+        max_frequency: float,
+        dynamic_range: float,
+        n_slices: int,
+    ) -> Image.Image:
+        ...

ultratrace2/model/files/loaders/parselmouth_spectrogram.py

@@ -0,0 +1,101 @@
+from PIL import Image  # type: ignore
+
+import logging
+import numpy as np
+import parselmouth  # type: ignore
+import pickle
+import warnings
+
+from .base import FileLoadError, SoundFileLoader, SpectrogramFileLoader
+
+
+logger = logging.getLogger(__name__)
+
+
+class ParselmouthLoader(SpectrogramFileLoader):
+    def get_path(self) -> str:
+        return self._path
+
+    def set_path(self, path) -> None:
+        self._path = path
+
+    def __init__(self, path: str, pm_sound: parselmouth.Sound):
+        self.set_path(path)
+        self.pm_sound = pm_sound
+
+    def save(self) -> None:
+        with open(self.get_path(), "wb") as fp:
+            pickle.dump(self.pm_sound.as_array(), fp)
+            logger.debug(f"Wrote ParselmouthLoader to disk: {self.get_path()}")
+
+    @classmethod
+    def from_file(cls, path: str) -> "ParselmouthLoader":
+        try:
+            with open(path, "rb") as fp:
+                pm_sound = parselmouth.Sound(pickle.load(fp))
+            return cls(path, pm_sound)
+        except Exception as e:
+            raise FileLoadError from e
+
+    load = from_file  # alias
+
+    @classmethod
+    def from_sound_file(cls, sound_file: SoundFileLoader) -> "ParselmouthLoader":
+        try:
+            pm_sound = parselmouth.Sound(sound_file.get_path())
+            pm_path = f"{sound_file.get_path()}.pmpkl"
+            pm = cls(pm_path, pm_sound)
+            pm.save()
+            return pm
+        except Exception as e:
+            raise FileLoadError from e
+
+    def get_image(
+        self,
+        start_time_ms: int,
+        stop_time_ms: int,
+        window_length: float,  # 0.005,
+        max_frequency: float,  # 5000,
+        dynamic_range: float,  # 90,
+        n_slices: int,  # 10 ** 7,
+    ) -> Image.Image:
+
+        # NB: need to do validation of *_time_ms BEFORE calling
+        start_time_sec = start_time_ms / 1000
+        stop_time_sec = stop_time_ms / 1000
+        visible_duration_sec = stop_time_sec - start_time_sec
+
+        time_step = visible_duration_sec / n_slices
+
+        spec = self.pm_sound.extract_part(
+            from_time=start_time_sec, to_time=stop_time_sec
+        ).to_spectrogram(
+            window_length=window_length,
+            maximum_frequency=max_frequency,
+            time_step=time_step,
+        )
+
+        with warnings.catch_warnings():
+            # We filter our numpy's warnings about division by zero, since we expect
+            # there to be some "padding" at the start and end of the recording with
+            # decibel-level zero.  Also, numpy does the sensible thing of defining
+            # `np.log10(0) := float("-inf")`, so it's not an issue for us.
+            warnings.filterwarnings(
+                "ignore",
+                message="divide by zero encountered in log10",
+                category=RuntimeWarning,
+            )
+            spec_arr = 10 * np.log10(np.flipud(spec.values))
+
+        # Normalize output array
+        max_magnitude = spec_arr.max()
+        spec_arr = (
+            (
+                spec_arr.clip(max_magnitude - dynamic_range, max_magnitude)
+                - max_magnitude
+            )
+            * -255
+            / dynamic_range
+        )
+
+        return Image.fromarray(spec_arr).convert("RGB")

ultratrace2/model/files/loaders/pydub.py

@@ -1,8 +1,12 @@
+import logging
 import pydub  # type: ignore
 
 from .base import FileLoadError, SoundFileLoader
 
 
+logger = logging.getLogger(__name__)
+
+
 class PydubLoader(SoundFileLoader):
     def get_path(self) -> str:
         return self._path

ultratrace2/model/files/loaders/tests/test_parselmouth_spectrogram.py

@@ -0,0 +1,78 @@
+import os
+import PIL  # type: ignore
+import pytest
+
+from ..base import SoundFileLoader
+from ..mp3 import MP3Loader
+from ..parselmouth_spectrogram import ParselmouthLoader
+
+
+@pytest.mark.parametrize(
+    "sound_file_loader, path",
+    [
+        (MP3Loader, "./test-data/example-audio/cello82.mp3"),
+        (MP3Loader, "./test-data/example-audio/harpsi-cs.mp3"),
+        (MP3Loader, "./test-data/example-audio/gtr-nylon22.mp3"),
+        (MP3Loader, "./test-data/example-audio/pno-cs.mp3"),
+        (MP3Loader, "./test-data/example-audio/bachfugue.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex016/link00.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex003/file01.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex004/file00.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex002/file00.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex015/file00.mp3"),
+        (
+            MP3Loader,
+            "./test-data/example-bundles/ex012/sub01/sub00/sub00/sub00/file00.mp3",
+        ),
+        (MP3Loader, "./test-data/example-bundles/ex009/file00.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex008/file01.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex008/file00.mp3"),
+        (MP3Loader, "./test-data/example-bundles/ex008/file02.mp3"),
+    ],
+)
+def test_loading_from_sound_file(sound_file_loader: SoundFileLoader, path: str) -> None:
+    try:
+        expected_pm_path = path + ".pmpkl"
+        assert not os.path.exists(expected_pm_path)
+        sound_file = sound_file_loader.from_file(path)
+        pm_file = ParselmouthLoader.from_sound_file(sound_file)
+        assert isinstance(pm_file, ParselmouthLoader)
+        assert pm_file.get_path() == expected_pm_path
+        assert os.path.exists(expected_pm_path)
+    finally:
+        if os.path.exists(expected_pm_path):
+            os.remove(expected_pm_path)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "./test-data/example-parselmouth/00.pmpkl",
+        "./test-data/example-parselmouth/01.pmpkl",
+        "./test-data/example-parselmouth/02.pmpkl",
+        "./test-data/example-parselmouth/03.pmpkl",
+        "./test-data/example-parselmouth/04.pmpkl",
+        "./test-data/example-parselmouth/05.pmpkl",
+        "./test-data/example-parselmouth/06.pmpkl",
+        "./test-data/example-parselmouth/07.pmpkl",
+        "./test-data/example-parselmouth/08.pmpkl",
+        "./test-data/example-parselmouth/09.pmpkl",
+        "./test-data/example-parselmouth/10.pmpkl",
+        "./test-data/example-parselmouth/11.pmpkl",
+        "./test-data/example-parselmouth/12.pmpkl",
+        "./test-data/example-parselmouth/13.pmpkl",
+        "./test-data/example-parselmouth/14.pmpkl",
+    ],
+)
+def test_loading_from_pmpkl_file(path: str) -> None:
+    pm_file = ParselmouthLoader.from_file(path)
+    assert isinstance(pm_file, ParselmouthLoader)
+    im = pm_file.get_image(
+        start_time_ms=0,
+        stop_time_ms=10,
+        window_length=0.005,
+        max_frequency=5000,
+        dynamic_range=90,
+        n_slices=10 ** 7,
+    )
+    assert isinstance(im, PIL.Image.Image)

ultratrace2/model/files/tests/test_bundle.py

@@ -19,7 +19,7 @@
 def test_empty_file_bundle_constructor(kwargs) -> None:
     fb = FileBundle("test", **kwargs)
     assert not fb.has_impl()
-    assert str(fb) == 'Bundle("test",None,None,None)'
+    assert str(fb) == 'Bundle("test",None,None,None,None)'
 
 
 def test_build_from_nonexistent_dir(mocker) -> None: