diff --git a/ultratrace2/model/files/__init__.py b/ultratrace2/model/files/__init__.py index bf613f7..31ef39b 100644 --- a/ultratrace2/model/files/__init__.py +++ b/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) diff --git a/ultratrace2/model/files/bundle.py b/ultratrace2/model/files/bundle.py index 7bf73ae..aff4f8b 100644 --- a/ultratrace2/model/files/bundle.py +++ b/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( diff --git a/ultratrace2/model/files/loaders/__init__.py b/ultratrace2/model/files/loaders/__init__.py index c27a3dd..6bc4171 100644 --- a/ultratrace2/model/files/loaders/__init__.py +++ b/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 diff --git a/ultratrace2/model/files/loaders/base.py b/ultratrace2/model/files/loaders/base.py index 99b4b6e..cc0ce44 100644 --- a/ultratrace2/model/files/loaders/base.py +++ b/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: + ... diff --git a/ultratrace2/model/files/loaders/parselmouth_spectrogram.py b/ultratrace2/model/files/loaders/parselmouth_spectrogram.py new file mode 100644 index 0000000..305919f --- /dev/null +++ b/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") diff --git a/ultratrace2/model/files/loaders/pydub.py b/ultratrace2/model/files/loaders/pydub.py index 470878c..50dba62 100644 --- a/ultratrace2/model/files/loaders/pydub.py +++ b/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 diff --git a/ultratrace2/model/files/loaders/tests/test_parselmouth_spectrogram.py b/ultratrace2/model/files/loaders/tests/test_parselmouth_spectrogram.py new file mode 100644 index 0000000..9963394 --- /dev/null +++ b/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) diff --git a/ultratrace2/model/files/tests/test_bundle.py b/ultratrace2/model/files/tests/test_bundle.py index fb0e558..d633473 100644 --- a/ultratrace2/model/files/tests/test_bundle.py +++ b/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: