amine@347: import os
amine@348: import sys
amine@403: from tempfile import TemporaryDirectory
amine@403: 
amine@403: import matplotlib
amine@400: import pytest
amine@348: 
amine@400: matplotlib.use("AGG")
amine@400: import matplotlib.pyplot as plt  # noqa E402
amine@403: 
amine@400: from auditok.core import AudioRegion  # noqa E402
amine@347: 
amine@400: SAVE_NEW_IMAGES = False
amine@400: if SAVE_NEW_IMAGES:
amine@400:     import shutil  # noqa E402
amine@400: 
amine@347: matplotlib.rcParams["figure.figsize"] = (10, 4)
amine@347: 
amine@347: 
amine@400: @pytest.mark.parametrize("channels", [1, 2], ids=["mono", "stereo"])
amine@400: def test_region_plot(channels):
amine@400:     type_ = "mono" if channels == 1 else "stereo"
amine@403:     audio_filename = "tests/data/test_split_10HZ_{}.raw".format(type_)
amine@426:     image_filename = "tests/images/plot_{}_region.png".format(type_)
amine@400:     expected_image = plt.imread(image_filename)
amine@400:     with TemporaryDirectory() as tmpdir:
amine@400:         output_image_filename = os.path.join(tmpdir, "image.png")
amine@400:         region = AudioRegion.load(audio_filename, sr=10, sw=2, ch=channels)
amine@400:         region.plot(show=False, save_as=output_image_filename)
amine@400:         output_image = plt.imread(output_image_filename)
amine@400: 
amine@400:         if SAVE_NEW_IMAGES:
amine@400:             shutil.copy(output_image_filename, image_filename)
amine@400:     assert (output_image == expected_image).all()  # mono, stereo
amine@400: 
amine@400: 
amine@400: @pytest.mark.parametrize(
amine@400:     "channels, use_channel",
amine@400:     [
amine@400:         (1, None),  # mono
amine@400:         (2, "any"),  # stereo_any
amine@400:         (2, 0),  # stereo_uc_0
amine@400:         (2, 1),  # stereo_uc_1
amine@400:         (2, "mix"),  # stereo_uc_mix
amine@400:     ],
amine@400:     ids=["mono", "stereo_any", "stereo_uc_0", "stereo_uc_1", "stereo_uc_mix"],
amine@400: )
amine@400: def test_region_split_and_plot(channels, use_channel):
amine@400:     type_ = "mono" if channels == 1 else "stereo"
amine@403:     audio_filename = "tests/data/test_split_10HZ_{}.raw".format(type_)
amine@400:     if type_ == "mono":
amine@426:         image_filename = "tests/images/split_and_plot_mono_region.png"
amine@400:     else:
amine@426:         image_filename = (
amine@426:             f"tests/images/split_and_plot_uc_{use_channel}_stereo_region.png"
amine@426:         )
amine@400: 
amine@400:     expected_image = plt.imread(image_filename)
amine@400:     with TemporaryDirectory() as tmpdir:
amine@400:         output_image_filename = os.path.join(tmpdir, "image.png")
amine@400:         region = AudioRegion.load(audio_filename, sr=10, sw=2, ch=channels)
amine@400:         region.split_and_plot(
amine@400:             aw=0.1,
amine@400:             uc=use_channel,
amine@400:             max_silence=0,
amine@400:             show=False,
amine@400:             save_as=output_image_filename,
amine@348:         )
amine@400:         output_image = plt.imread(output_image_filename)
amine@347: 
amine@400:         if SAVE_NEW_IMAGES:
amine@400:             shutil.copy(output_image_filename, image_filename)
amine@400:     assert (output_image == expected_image).all()