ambiviz 0.0.1__tar.gz

ambiviz-0.0.1/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: ambiviz
+Version: 0.0.1
+Summary: Visualization tools for ambisonic audio and 360 video
+Author-email: Arthur Jinyue Guo <jinyue.guo@imv.uio.no>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=2.0.2
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: librosa>=0.10
+Requires-Dist: torchaudio>=2.0
+Requires-Dist: opencv-python>=4.8
+Requires-Dist: Pillow>=10.0
+Requires-Dist: tqdm>=4.0
+Requires-Dist: scipy>=1.10
+Requires-Dist: scienceplots>=2.2.1
+Requires-Dist: musicalgestures>=1.3.3
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+# ambiviz
+
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+[![Actions status](https://github.com/fisheggg/ambiviz/workflows/ci.yml/badge.svg)](https://github.com/fisheggg/ambiviz/actions)
+
+
+Visualization tools for ambisonic audio (and 360 video).
+
+
+## References
+
+- The core ambisonics decoder is modified from [https://github.com/pedro-morgado/spatialaudiogen](https://github.com/pedro-morgado/spatialaudiogen).

ambiviz-0.0.1/README.md

@@ -0,0 +1,12 @@
+# ambiviz
+
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+[![Actions status](https://github.com/fisheggg/ambiviz/workflows/ci.yml/badge.svg)](https://github.com/fisheggg/ambiviz/actions)
+
+
+Visualization tools for ambisonic audio (and 360 video).
+
+
+## References
+
+- The core ambisonics decoder is modified from [https://github.com/pedro-morgado/spatialaudiogen](https://github.com/pedro-morgado/spatialaudiogen).

ambiviz-0.0.1/ambiviz/__init__.py

@@ -0,0 +1,10 @@
+from .aem import aem
+from .anglegram import anglegram
+from .directogram import directogram
+
+
+__all__ = [
+    "aem",
+    "anglegram",
+    "directogram",
+]

ambiviz-0.0.1/ambiviz/aem.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import glob
+import os
+import subprocess
+import logging
+from pathlib import Path
+from typing import List, Optional
+
+import cv2
+import librosa
+import musicalgestures
+import numpy as np
+import PIL
+import tqdm
+
+from .ambisonics.audio_to_aem import compute_aem
+
+logger = logging.getLogger(__name__)
+
+
+def aem(
+    audio_path: str,
+    video_paths: List[str],
+    save_path: str | Path,
+    duration: float,
+    audio_offset: float = 0.0,
+    video_offset: float = 0.0,
+    normalization_mode: str = "global",
+    use_gpu: bool = False,
+    aem_mode: str = "aem",
+    aem_mask_alpha: float = 0.5,
+    verbose: bool = False,
+    to_db: bool = False,
+    **kwargs,
+):
+    """
+    Takes an FOA audio and a 360 video, plots the audio energy map on the video.
+
+    Parameters
+    ----------
+    normalization_mode: "global" or "local"
+        if "global", normalize the AEM according to the min/max over the entire video.
+        if "local", normalize the AEM of each frame by itself
+    """
+    assert normalization_mode in ["global", "local"]
+    assert aem_mode in ["aem"]
+
+    if verbose:
+        logging.basicConfig(level=logging.INFO)
+
+    # Load video
+    logger.info("Loading video...")
+    video, fps = musicalgestures.Mg360Video(
+        video_paths,
+        projection="equirectangular",
+        starttime=video_offset,
+        endtime=video_offset + duration,
+    ).numpy()
+    video = np.copy(video)  # make writable
+    logger.info(f"Video shape: {video.shape}")
+    logger.info(f"Video FPS: {fps}")
+    logger.info(
+        f"Video stats: max={np.max(video)}, min={np.min(video)}, mean={np.mean(video)}"
+    )
+    logger.info(f"Video dtype: {video.dtype}")
+
+    video_frame_shape = video.shape[1:3]
+    video = video * 0.95  # reduce brightness
+
+    audio_sr = librosa.get_samplerate(audio_path)
+    aem, timestamp, phi_mesh, nu_mesh = compute_aem(
+        audio_path,
+        mode=aem_mode,
+        fps=fps,
+        audio_frame_length=int(audio_sr * 4 / fps),
+        audio_hop_length=int(audio_sr * 2 / fps),
+        gpu=use_gpu,
+        verbose=verbose,
+        duration=duration,
+        offset=audio_offset,
+        to_db=to_db,
+        **kwargs,
+    )
+    # cv2 counts video frame from top left,
+    # aem counts from bottom left, so flip the aem vertically
+    aem = np.flip(aem, axis=1)
+    aem = np.repeat(
+        aem, 2, axis=0
+    )  # when hoplength reduced to half fps, repeat each AEM frame once to match the length
+
+    if normalization_mode == "global":
+        # global normalize
+        aem = (aem - np.min(aem)) / (np.max(aem) - np.min(aem))
+        # aem_frame_gain = np.mean(aem, axis=(1, 2))
+        # aem_frame_gain = (aem_frame_gain - np.min(aem_frame_gain)) / (
+        #     np.max(aem_frame_gain) - np.min(aem_frame_gain)
+        # )
+    logger.info(f"AEM shape: {aem.shape}")
+    logger.info(f"AEM stats: max={np.max(aem)}, min={np.min(aem)}, mean={np.mean(aem)}")
+
+    # add aem colormask to video
+    frames = range(min(video.shape[0], aem.shape[0]))
+    if verbose:
+        frames = tqdm.tqdm(frames, desc="=> rendering AEM mask to video")
+    for i in frames:
+        # aem_frame = aem[i]
+        # aem_frame = aem_frame_gain[i] * (aem[i] - np.min(aem[i])) / (np.max(aem[i]) - np.min(aem[i])) # gain+local normalization
+        if normalization_mode == "local":
+            # local normalize
+            aem_frame = (aem[i] - np.min(aem[i])) / (np.max(aem[i]) - np.min(aem[i]))
+        elif normalization_mode == "global":
+            # make local changes more drastic, but keeps the local max of the frame
+            aem_frame = aem[i]
+            aem_frame_max = np.max(aem_frame)
+            aem_frame = aem_frame**8
+            aem_frame = aem_frame / np.max(aem_frame) * aem_frame_max
+
+        aem_frame = (aem_mask_alpha * aem_frame * 255).astype(video.dtype)
+        # aem_frame = (aem_mask_alpha * aem[i] * 255).astype(video.dtype)
+        aem_frame_resized = PIL.Image.fromarray(aem_frame).resize(
+            (video_frame_shape[1], video_frame_shape[0])
+        )
+        # video[i, :, :, 2] = (video[i, :, :, 2] * (1 - aem_mask_alpha) + np.asarray(aem_frame_resized)).astype(video.dtype)
+        video[i, :, :, 2] = (video[i, :, :, 2] + np.asarray(aem_frame_resized)).astype(
+            video.dtype
+        )
+    video = np.clip(video, 0, 255)
+
+    # save video
+    save_path = Path(save_path).resolve().absolute()
+    audio_path = Path(audio_path).resolve().absolute()
+    if not save_path.suffix:  # add a filename if a dir is given
+        save_path = save_path / f"{Path(audio_path).stem}.mp4"
+        save_path.mkdir(parents=True, exist_ok=True)
+    else:
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+    save_path_tmp = save_path.with_stem(save_path.stem + "_tmp")
+    video = video.astype(np.uint8)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(str(save_path_tmp), fourcc, fps, video_frame_shape[::-1])
+    range_ = range(video.shape[0])
+    if verbose:
+        range_ = tqdm.tqdm(range_, desc="=> writing video...")
+    for frame_idx in range_:
+        out.write(video[frame_idx])
+    out.release()
+
+    # merge audio and video, only take the first channel from ambisonics
+    cmd = f"ffmpeg -hide_banner -y -i '{save_path_tmp}' -ss {audio_offset} -i '{audio_path}' -filter_complex \"[1:a:0]pan=mono|c0=c0[a_out]\" -map 0:v:0 -map \"[a_out]\" -c:v copy -c:a aac -shortest '{save_path}'"
+    # cmd = [
+    #     "ffmpeg",
+    #     "-hide_banner",
+    #     "-y",
+    #     "-i",
+    #     f"{save_path_tmp}",
+    #     "-ss",
+    #     f"{audio_offset}",
+    #     "-i",
+    #     f"{audio_path}",
+    #     "-filter_complex",
+    #     "[1:a:0]pan=mono|c0=c0[a_out]",
+    #     "-map",
+    #     "[a_out]",
+    #     "-c:v",
+    #     "copy",
+    #     "-c:a",
+    #     "aac",
+    #     "-shortest",
+    #     f"{save_path}",
+    # ]
+    subprocess.call(cmd, shell=True)
+    # subprocess.check_call(cmd)
+    os.remove(save_path_tmp)
+
+    if verbose:
+        print(f"=> Video saved to {save_path}")
+
+
+if __name__ == "__main__":
+    audio_path = "../test_files/test_aem_3oa.wav"
+    video_path = "../test_files/test_aem_video.mp4"
+
+    aem(
+        audio_path,
+        [video_path],
+        "../test_output",
+        duration=30,
+        audio_offset=23.3 + 270,
+        video_offset=32.5 + 270,
+        verbose=True,
+        aem_mask_alpha=0.5,
+        ambi_order=3,
+        to_db=True,
+    )

ambiviz-0.0.1/ambiviz/anglegram.py

@@ -0,0 +1,275 @@
+import datetime
+import logging
+import os
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import librosa
+import matplotlib.pyplot as plt
+import numpy as np
+import scienceplots
+
+from .ambisonics.audio_to_aem import compute_aem
+
+logger = logging.getLogger(__name__)
+
+
+def compute_anglegram_from_aem(
+    aem_path: Optional[Union[str, os.PathLike, Path]] = None,
+    time_stamp=None,
+    phi_mesh=None,
+    nu_mesh=None,
+    aem=None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    if aem_path is not None:
+        assert (
+            time_stamp is None and phi_mesh is None and nu_mesh is None and aem is None
+        ), "If aem_path is provided, other parameters should be None."
+    else:
+        assert (
+            time_stamp is not None
+            and phi_mesh is not None
+            and nu_mesh is not None
+            and aem is not None
+        ), "If aem_path is not provided, other parameters should not be None."
+    if aem_path is not None:
+        data = np.load(aem_path)
+        time_stamp, phi_mesh, nu_mesh, aem = (
+            data["time_stamp"],
+            data["phi_mesh"],
+            data["nu_mesh"],
+            data["aem"],
+        )
+    phi_record, nu_record, energy_record = [], [], []
+
+    aem_argmax = [np.unravel_index(np.argmax(r), r.shape) for r in aem]
+    for i, point in enumerate(aem_argmax):
+        phi_record.append(phi_mesh[point[0], point[1]])
+        nu_record.append(nu_mesh[point[0], point[1]])
+        energy_record.append(aem[i, point[0], point[1]])
+    phi_record = np.array(phi_record)
+    nu_record = np.array(nu_record)
+    energy_record = np.array(energy_record)
+    energy_record = librosa.amplitude_to_db(energy_record, ref=np.max)
+
+    return time_stamp, phi_record, nu_record, energy_record
+
+
+# def get_anglegram_mean_rms(audio_path) -> Tuple[float, np.ndarray]:
+#     """
+#     Calculate average rms of the audio file, made for noise estimation..
+#     """
+#     time_stamp, phi_mesh, nu_mesh, aem = compute_aem(
+#         audio_path=audio_path,
+#         fps=20,
+#         audio_frame_length=4800,
+#         aem_width=360,  # angular_res = 1
+#         aem_height=180,
+#         gpu=True,
+#         batch_size=5,
+#     )
+#     phi_record, nu_record, energy_record = [], [], []
+
+#     aem_argmax = [np.unravel_index(np.argmax(r), r.shape) for r in aem]
+#     for i, point in enumerate(aem_argmax):
+#         phi_record.append(phi_mesh[point[0], point[1]])
+#         nu_record.append(nu_mesh[point[0], point[1]])
+#         energy_record.append(aem[i, point[0], point[1]])
+#     phi_record = np.array(phi_record)
+#     nu_record = np.array(nu_record)
+#     energy_record = np.array(energy_record)
+
+#     return np.mean(energy_record), energy_record
+
+
+def plot_anglegram(
+    time_record,
+    phi_record,
+    nu_record,
+    rms_record,
+    db_threshold=-60,
+    db_max=-20,
+    title: Optional[str] = None,
+    save_path: Optional[Union[str, os.PathLike[str], Path]] = None,
+):
+    with plt.style.context(["science", "no-latex"]):
+        fig, ax = plt.subplots(2, 1, figsize=(12, 6))
+        if title:
+            fig.suptitle(title, fontsize=16)
+
+        sc = ax[0].scatter(
+            time_record[rms_record > db_threshold],
+            phi_record[rms_record > db_threshold],
+            c=rms_record[rms_record > db_threshold],
+            label="Azimuth",
+            alpha=0.9,
+            vmin=db_threshold,
+            vmax=db_max,
+        )
+        # find the closest multiple of 2pi to azimuth_min
+        azimuth_min_pi = np.floor(np.min(phi_record) / np.pi) * np.pi
+        azimuth_max_pi = np.ceil(np.max(phi_record) / np.pi) * np.pi
+        ax[0].set_ylim(azimuth_min_pi - 0.1, azimuth_max_pi + 0.1)
+        yticks = np.arange(azimuth_min_pi, azimuth_max_pi + 0.1, np.pi)
+        ax[0].set_yticks(yticks)
+        ax[0].set_yticklabels([f"{int(ytick / np.pi)}$\pi$" for ytick in yticks])
+
+        # plt.colorbar(sc, ax=ax[0], label="RMS (dB)")
+        ax[0].set(xlabel="Time (s)", ylabel="Azimuth (rad)")
+        # horizontal lines on pi and 2pi
+        for ytick in yticks:
+            if np.abs(ytick % (2 * np.pi)) < 1e-1:
+                ax[0].axhline(y=ytick, linestyle="-", color="black", alpha=0.5)
+            else:
+                ax[0].axhline(y=ytick, linestyle="--", alpha=0.3)
+
+        ax[0].grid(axis="x")
+
+        sc = ax[1].scatter(
+            time_record[rms_record > db_threshold],
+            nu_record[rms_record > db_threshold],
+            c=rms_record[rms_record > db_threshold],
+            label="Elevation",
+            alpha=0.9,
+            vmin=db_threshold,
+            vmax=db_max,
+        )
+        # ax[1].plot(time_record, nu_record, label='Elevation')
+        ax[1].set(
+            xlabel="Time (s)",
+            ylabel="Elevation (rad)",
+            ylim=(-0.5 * np.pi - 0.1, 0.5 * np.pi + 0.1),
+        )
+        ax[1].grid(axis="x")
+        yticks = np.array([-0.5 * np.pi, -0.25 * np.pi, 0, 0.25 * np.pi, 0.5 * np.pi])
+        ax[1].set_yticks(yticks)
+        ax[1].set_yticklabels(
+            [r"$-\pi/2$", r"$-\pi/4$", r"$0$", r"$\pi/4$", r"$\pi/2$"]
+        )
+        for ytick in yticks:
+            if ytick == 0:
+                ax[1].axhline(y=ytick, linestyle="-", color="black", alpha=0.5)
+            else:
+                ax[1].axhline(y=ytick, linestyle="--", alpha=0.3)
+        # plt.colorbar(sc, ax=ax[1], label="RMS (dB)")
+        fig.colorbar(sc, ax=ax.ravel().tolist(), label="RMS (dB)")
+
+    if save_path is not None:
+        plt.savefig(save_path)
+        logger.info(f"Figure saved to {save_path}")
+
+    return (
+        ax,
+        time_record[rms_record > db_threshold],
+        phi_record[rms_record > db_threshold],
+        nu_record[rms_record > db_threshold],
+        rms_record[rms_record > db_threshold],
+    )
+
+
+def plot_anglemap(aem_path: str, save_path: Optional[str] = None, **kwargs):
+    data = np.load(aem_path)
+    time_stamp, _, _, aem = (
+        data["time_stamp"],
+        data["phi_mesh"],
+        data["nu_mesh"],
+        data["aem"],
+    )
+
+    logger.debug(f"aem shape: {aem.shape}")  # (1093, 180, 360)
+
+    phi_img = np.mean(aem, axis=1)  # (1093, 360)
+    phi_img = np.flip(phi_img, axis=1).T  # (360, 1093), top left pi
+    logger.debug(f"phi_img shape: {phi_img.shape}")
+
+    nu_img = np.mean(aem, axis=2)  # (1093, 180)
+    nu_img = np.flip(nu_img, axis=1).T  # (180, 1093), top left pi/2
+    logger.debug(f"nu_img shape: {nu_img.shape}")
+
+    fig, ax = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
+    fig.suptitle(
+        f"{os.path.basename(aem_path)}\nmin: {aem.mean():.2f}, max: {aem.max():.2f}",
+        fontsize=16,
+    )
+    ax[0].imshow(phi_img, cmap="hot", aspect="auto", **kwargs)
+    ax[0].set_yticks([0, 89, 179, 269, 359])
+    ax[0].set_yticklabels([r"$\pi$", r"$\pi/2$", r"$0$", r"$-\pi/2$", r"$-\pi$"])
+    ax[0].set_xticks(np.arange(0, len(time_stamp), 200))
+    ax[0].set_xticklabels(np.arange(0, len(time_stamp), 200) / 20)
+
+    ax[0].set_ylabel("Azimuth (rad)")
+
+    ax[1].imshow(nu_img, cmap="hot", aspect="auto", **kwargs)
+    ax[1].set_xlabel("Time (s)")
+    ax[1].set_ylabel("Elevation (rad)")
+    ax[1].set_yticks([0, 44, 89, 134, 179])
+    ax[1].set_yticklabels([r"$\pi/2$", r"$\pi/4$", r"$0$", r"$-\pi/4$", r"$-\pi/2$"])
+
+    if save_path is not None:
+        plt.savefig(save_path)
+        logger.info(f"=> Figure saved to {save_path}")
+    else:
+        plt.show()
+
+
+def frame_to_seconds(frame_idx, fps=30):
+    return frame_idx / fps
+
+
+def seconds_to_timestring(seconds):
+    return str(datetime.timedelta(seconds=seconds))
+
+
+def frame_to_timestring(frame_idx, fps=30):
+    return seconds_to_timestring(frame_to_seconds(frame_idx, fps))
+
+
+def anglegram(
+    audio_path: Union[str, os.PathLike[str], Path],
+    save_path: Optional[Union[str, os.PathLike[str], Path]] = None,
+    fps: int = 20,
+    audio_frame_length: int = 4800,
+    aem_width: int = 90,
+    aem_height: int = 45,
+    gpu: bool = False,
+    db_threshold: int = -60,
+    db_max: int = -20,
+    title: Optional[str] = None,
+    batch_size: int = 5,
+):
+    logger.debug(f"Computing AEM. Audio path: {audio_path}")
+    aem, time_stamp, phi_mesh, nu_mesh = compute_aem(
+        audio_path=audio_path,
+        fps=fps,
+        audio_frame_length=audio_frame_length,
+        aem_width=aem_width,
+        aem_height=aem_height,
+        gpu=gpu,
+        batch_size=batch_size,
+    )
+
+    logger.debug(f"Computing anglegram. AEM shape: {aem.shape}")
+    timestamp, phi_record, nu_record, rms_record = compute_anglegram_from_aem(
+        time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh, aem=aem
+    )
+    logger.debug(
+        f"rms record max: {rms_record.max()}, min: {rms_record.min()}, mean: {rms_record.mean()}"
+    )
+
+    logger.debug(f"Plotting anglegram. Save path: {save_path}")
+    if title is None:
+        title = os.path.basename(audio_path)
+
+    ax, *_ = plot_anglegram(
+        timestamp,
+        phi_record,
+        nu_record,
+        rms_record,
+        db_threshold=db_threshold,
+        db_max=db_max,
+        title=title,
+        save_path=save_path,
+    )
+
+    if not save_path:
+        return ax

ambiviz-0.0.1/ambiviz/directogram.py

@@ -0,0 +1,98 @@
+import logging
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import scienceplots
+
+from .ambisonics.audio_to_aem import compute_aem
+
+logger = logging.getLogger(__name__)
+
+
+def plot_directogram(
+    angles,
+    energy,
+    title=None,
+    figsize=(8, 8),
+    dpi=300,
+    ax=None,
+    save_path=None,
+):
+    with plt.style.context(["science", "no-latex"]):
+        if ax is None:
+            fig = plt.figure(figsize=figsize, dpi=dpi)
+            ax = fig.add_subplot(111, projection="polar")
+            ax.set_title(title if title is not None else "Directogram")
+        ax.bar(
+            x=angles / 180 * np.pi,
+            height=energy,
+            width=2 * np.pi / len(angles),
+            color="tab:blue",
+            alpha=0.7,
+            label="Summed directional energy",
+        )
+        ax.set_theta_zero_location("N")
+
+    if save_path is not None:
+        plt.savefig(save_path)
+        logger.info(f"Figure saved to {save_path}")
+
+    return ax
+
+
+def compute_directogram_from_aem(aem, time_stamp, phi_mesh, nu_mesh):
+    """
+    Calculate the directogram from the AEM.
+
+    Args:
+        aem: np.ndarray, shape (n_frames, n_nu, n_phi)
+        time_stamp: np.ndarray, shape (n_frames,)
+        phi_mesh: np.ndarray, azimuth. shape: (n_nu, n_phi)
+        nu_mesh: np.ndarray, elevation. shape: (n_nu, n_phi)
+
+    Returns:
+        angles: np.ndarray, in degrees, shape (n_phi,)
+        energy: np.ndarray, shape (n_phi,)
+    """
+    horizontal_nu_idx = nu_mesh.shape[0] // 2
+    horizontal_aem = aem[:, horizontal_nu_idx, :]
+    angles = phi_mesh[0, :] / np.pi * 180
+    energy = np.zeros_like(angles)
+    for i in range(horizontal_aem.shape[0]):
+        argmax_idx = np.argmax(horizontal_aem[i])
+        energy[argmax_idx] += horizontal_aem[i, argmax_idx]
+
+    return angles, energy
+
+
+def directogram(
+    audio_path: Union[str, os.PathLike, Path],
+    save_path: Optional[Union[str, os.PathLike, Path]] = None,
+    ax=None,
+    fps: int = 20,
+    audio_frame_length: int = 4800,
+    aem_width: int = 90,
+    aem_height: int = 45,
+    batch_size: int = 5,
+    gpu: bool = False,
+):
+    aem, time_stamp, phi_mesh, nu_mesh = compute_aem(
+        audio_path=audio_path,
+        fps=fps,
+        audio_frame_length=audio_frame_length,
+        aem_width=aem_width,
+        aem_height=aem_height,
+        gpu=gpu,
+        batch_size=batch_size,
+    )
+
+    logger.debug(f"Computing anglegram. AEM shape: {aem.shape}")
+    angles, energy = compute_directogram_from_aem(aem, time_stamp, phi_mesh, nu_mesh)
+
+    logger.debug(f"Plotting directogram. Save path: {save_path}")
+    ax = plot_directogram(angles, energy, ax=ax, save_path=save_path)
+
+    return ax, angles, energy

ambiviz-0.0.1/ambiviz/utils.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+
+def check_ambisonics_order(audio: np.ndarray) -> int:
+    """
+    Check the order of ambisonics.
+    Raises error if the number of channels does not correspond to a valid ambisonics order.
+
+    Parameters:
+        audio: (n_channels, n_samples)
+
+    Returns:
+        order: int, the order of ambisonics
+    """
+    assert len(audio.shape) == 2, "Audio should be 2D (n_channels, n_samples)"
+    n_channels = audio.shape[0]
+
+    order = int(np.sqrt(n_channels) - 1)
+    if (order + 1) ** 2 != n_channels:
+        raise ValueError(
+            f"Number of channels {n_channels} does not correspond to a valid ambisonics order."
+        )
+
+    return order

ambiviz-0.0.1/ambiviz.egg-info/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: ambiviz
+Version: 0.0.1
+Summary: Visualization tools for ambisonic audio and 360 video
+Author-email: Arthur Jinyue Guo <jinyue.guo@imv.uio.no>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=2.0.2
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: librosa>=0.10
+Requires-Dist: torchaudio>=2.0
+Requires-Dist: opencv-python>=4.8
+Requires-Dist: Pillow>=10.0
+Requires-Dist: tqdm>=4.0
+Requires-Dist: scipy>=1.10
+Requires-Dist: scienceplots>=2.2.1
+Requires-Dist: musicalgestures>=1.3.3
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+# ambiviz
+
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+[![Actions status](https://github.com/fisheggg/ambiviz/workflows/ci.yml/badge.svg)](https://github.com/fisheggg/ambiviz/actions)
+
+
+Visualization tools for ambisonic audio (and 360 video).
+
+
+## References
+
+- The core ambisonics decoder is modified from [https://github.com/pedro-morgado/spatialaudiogen](https://github.com/pedro-morgado/spatialaudiogen).

ambiviz-0.0.1/ambiviz.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+README.md
+pyproject.toml
+ambiviz/__init__.py
+ambiviz/aem.py
+ambiviz/anglegram.py
+ambiviz/directogram.py
+ambiviz/utils.py
+ambiviz.egg-info/PKG-INFO
+ambiviz.egg-info/SOURCES.txt
+ambiviz.egg-info/dependency_links.txt
+ambiviz.egg-info/requires.txt
+ambiviz.egg-info/top_level.txt
+tests/test_aem.py
+tests/test_anglegram.py
+tests/test_directogram.py
+tests/test_plot.py

ambiviz-0.0.1/ambiviz.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

ambiviz-0.0.1/ambiviz.egg-info/requires.txt

@@ -0,0 +1,13 @@
+numpy>=2.0.2
+matplotlib>=3.7
+librosa>=0.10
+torchaudio>=2.0
+opencv-python>=4.8
+Pillow>=10.0
+tqdm>=4.0
+scipy>=1.10
+scienceplots>=2.2.1
+musicalgestures>=1.3.3
+
+[dev]
+pytest>=7.0

ambiviz-0.0.1/ambiviz.egg-info/top_level.txt

@@ -0,0 +1 @@
+ambiviz

ambiviz-0.0.1/pyproject.toml

@@ -0,0 +1,42 @@
+[project]
+name = "ambiviz"
+version = "0.0.1"
+description = "Visualization tools for ambisonic audio and 360 video"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [
+    {name = "Arthur Jinyue Guo", email = "jinyue.guo@imv.uio.no"},
+]
+dependencies = [
+    "numpy>=2.0.2",
+    "matplotlib>=3.7",
+    "librosa>=0.10",
+    "torchaudio>=2.0",
+    "opencv-python>=4.8",
+    "Pillow>=10.0",
+    "tqdm>=4.0",
+    "scipy>=1.10",
+    "scienceplots>=2.2.1",
+    "musicalgestures>=1.3.3",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0"]
+
+[tool.pytest.ini_options]
+markers = [
+    "integration: tests that require real audio/video files and ffmpeg (deselect with '-m not integration')",
+]
+
+[tool.ruff]
+ignore = ["F401"]
+
+[tool.setuptools]
+packages = ["ambiviz"]
+
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.31.0",
+    "pytest>=8.4.2",
+]

ambiviz-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

ambiviz-0.0.1/tests/test_aem.py

@@ -0,0 +1,38 @@
+"""Unit and integration tests for ambiviz.aem."""
+
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from ambiviz.aem import aem
+
+
+# Real test files (used by integration tests)
+TEST_FILES_DIR = Path(__file__).parent.parent / "test_files"
+REAL_AUDIO_PATH = TEST_FILES_DIR / "test_aem_3oa.wav"
+REAL_VIDEO_PATH = TEST_FILES_DIR / "test_aem_video.mp4"
+TEST_OUTPUT_DIR = Path(__file__).parent.parent / "test_output"
+
+
+@pytest.mark.integration
+def test_anglegram_integration():
+    """Integration test for aem"""
+
+    try:
+        aem(
+            audio_path=REAL_AUDIO_PATH,
+            video_paths=[str(REAL_VIDEO_PATH)],
+            save_path=str(TEST_OUTPUT_DIR / "test_aem_output.mp4"),
+            duration=30,
+            audio_offset=23.3 + 270,
+            video_offset=32.5 + 270,
+            aem_mask_alpha=0.5,
+            ambi_order=3,
+            to_db=True,
+            normalization_mode="local",
+            verbose=True,
+        )
+    except Exception as e:
+        pytest.fail(f"Integration test for aem failed with exception: {e}")

ambiviz-0.0.1/tests/test_anglegram.py

@@ -0,0 +1,131 @@
+"""Unit and integration tests for ambiviz.anglegram."""
+
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from ambiviz.anglegram import anglegram, compute_anglegram_from_aem
+
+# Real test files (used by integration tests)
+TEST_FILES_DIR = Path(__file__).parent.parent / "test_files"
+REAL_AUDIO_PATH = TEST_FILES_DIR / "test_anglegram.wav"
+TEST_OUTPUT_DIR = Path(__file__).parent.parent / "test_output"
+
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_aem_data(n_frames=10, height=4, width=8):
+    """Return fake (aem, time_stamp, phi_mesh, nu_mesh) as compute_aem would."""
+    rng = np.random.default_rng(0)
+    aem = rng.random((n_frames, height, width)).astype(np.float32)
+    time_stamp = np.linspace(0, 1, n_frames)
+    phi_mesh, nu_mesh = np.meshgrid(
+        np.linspace(-np.pi, np.pi, width),
+        np.linspace(-np.pi / 2, np.pi / 2, height),
+    )
+    return aem, time_stamp, phi_mesh, nu_mesh
+
+
+def _make_plot_return():
+    n = 5
+    return (MagicMock(), np.zeros(n), np.zeros(n), np.zeros(n), np.zeros(n))
+
+
+# ---------------------------------------------------------------------------
+# compute_anglegram_from_aem unit tests
+# ---------------------------------------------------------------------------
+
+
+class TestComputeAnglegramFromAem:
+    def test_output_shapes(self):
+        aem, time_stamp, phi_mesh, nu_mesh = _make_aem_data(n_frames=10)
+        ts, phi, nu, energy = compute_anglegram_from_aem(
+            time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh, aem=aem
+        )
+        assert ts.shape == (10,)
+        assert phi.shape == (10,)
+        assert nu.shape == (10,)
+        assert energy.shape == (10,)
+
+    def test_energy_is_db(self):
+        """Energy record should be in dB: max value must be 0.0."""
+        aem, time_stamp, phi_mesh, nu_mesh = _make_aem_data(n_frames=8)
+        _, _, _, energy = compute_anglegram_from_aem(
+            time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh, aem=aem
+        )
+        assert np.max(energy) == pytest.approx(0.0, abs=1e-5)
+
+    def test_energy_nonpositive(self):
+        """All dB values should be <= 0."""
+        aem, time_stamp, phi_mesh, nu_mesh = _make_aem_data(n_frames=8)
+        _, _, _, energy = compute_anglegram_from_aem(
+            time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh, aem=aem
+        )
+        assert np.all(energy <= 0.0)
+
+    def test_picks_argmax_coordinates(self):
+        """phi/nu should come from the grid point with the highest AEM value per frame."""
+        n_frames, height, width = 3, 4, 8
+        aem = np.zeros((n_frames, height, width), dtype=np.float32)
+        phi_mesh, nu_mesh = np.meshgrid(
+            np.linspace(-np.pi, np.pi, width),
+            np.linspace(-np.pi / 2, np.pi / 2, height),
+        )
+        time_stamp = np.arange(n_frames, dtype=float)
+
+        # Place the maximum at known positions
+        expected_phi = []
+        expected_nu = []
+        for i in range(n_frames):
+            row, col = i % height, (i * 2) % width
+            aem[i, row, col] = 1.0
+            expected_phi.append(phi_mesh[row, col])
+            expected_nu.append(nu_mesh[row, col])
+
+        _, phi, nu, _ = compute_anglegram_from_aem(
+            time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh, aem=aem
+        )
+        np.testing.assert_allclose(phi, expected_phi)
+        np.testing.assert_allclose(nu, expected_nu)
+
+    def test_requires_consistent_args(self):
+        """Passing aem_path together with other args should raise AssertionError."""
+        aem, time_stamp, phi_mesh, nu_mesh = _make_aem_data()
+        with pytest.raises(AssertionError):
+            compute_anglegram_from_aem(
+                aem_path="dummy.npz",
+                time_stamp=time_stamp,
+                phi_mesh=phi_mesh,
+                nu_mesh=nu_mesh,
+                aem=aem,
+            )
+
+    def test_requires_all_args_when_no_path(self):
+        """Omitting any required arg when aem_path is None should raise AssertionError."""
+        aem, time_stamp, phi_mesh, nu_mesh = _make_aem_data()
+        with pytest.raises(AssertionError):
+            compute_anglegram_from_aem(
+                time_stamp=time_stamp, phi_mesh=phi_mesh, nu_mesh=nu_mesh
+            )
+
+
+# ---------------------------------------------------------------------------
+# Integration test
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_anglegram_integration():
+    """Integration test for anglegram function."""
+    try:
+        anglegram(
+            audio_path=REAL_AUDIO_PATH,
+            save_path=TEST_OUTPUT_DIR / "test_anglegram.png",
+        )
+    except Exception as e:
+        pytest.fail(f"Integration test failed with exception: {e}")

ambiviz-0.0.1/tests/test_directogram.py

@@ -0,0 +1,33 @@
+"""Unit and integration tests for ambiviz.directogram."""
+
+import logging
+from pathlib import Path
+
+import pytest
+
+from ambiviz.directogram import directogram
+
+# Real test files (used by integration tests)
+TEST_FILES_DIR = Path(__file__).parent.parent / "test_files"
+REAL_AUDIO_PATH = TEST_FILES_DIR / "test_directogram.wav"
+TEST_OUTPUT_DIR = Path(__file__).parent.parent / "test_output"
+
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+# ---------------------------------------------------------------------------
+# Integration test
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_directogram_integration():
+    """Integration test for directogram function."""
+    try:
+        directogram(
+            audio_path=REAL_AUDIO_PATH,
+            save_path=TEST_OUTPUT_DIR / "test_directogram.png",
+        )
+    except Exception as e:
+        pytest.fail(f"Integration test failed with exception: {e}")

ambiviz-0.0.1/tests/test_plot.py

@@ -0,0 +1,117 @@
+# """Unit and integration tests for ambiviz.plot.visualize_aem."""
+# from pathlib import Path
+
+# import numpy as np
+# import pytest
+
+# from ambiviz.aem import visualize_aem
+
+# AUDIO_PATH = "fake/audio.wav"
+# VIDEO_PATH = "fake/video.mp4"
+
+# # Real test files (used by integration tests)
+# TEST_FILES_DIR = Path(__file__).parent.parent / "test_files"
+# REAL_AUDIO_PATH = TEST_FILES_DIR / "test_audio_foa.wav"
+# REAL_VIDEO_PATH = TEST_FILES_DIR / "test_video.mp4"
+
+
+# def _call(tmp_path, normalization_mode=None, **kwargs):
+#     """Helper: invoke visualize_aem with a concrete save path."""
+#     visualize_aem(
+#         AUDIO_PATH,
+#         VIDEO_PATH,
+#         str(tmp_path / "out.mp4"),
+#         duration=2.0,
+#         audio_offset=0.0,
+#         video_offset=0.0,
+#         normalization_mode=normalization_mode,
+#         **kwargs,
+#     )
+
+
+# # ---------------------------------------------------------------------------
+# # Validation
+# # ---------------------------------------------------------------------------
+
+# def test_invalid_normalization_mode_raises(tmp_path):
+#     with pytest.raises(AssertionError):
+#         _call(tmp_path, normalization_mode="bad_mode")
+
+
+# def test_invalid_aem_mode_raises(tmp_path):
+#     with pytest.raises(AssertionError):
+#         _call(tmp_path, aem_mode="bad_mode")
+
+
+# # ---------------------------------------------------------------------------
+# # Smoke tests — one per normalization mode
+# # ---------------------------------------------------------------------------
+
+# @pytest.mark.parametrize("normalization_mode", [None, "global", "local"])
+# def test_smoke(patched_env, normalization_mode):
+#     """Function completes without error for every normalization mode."""
+#     _call(patched_env["tmp_path"], normalization_mode=normalization_mode)
+
+
+# # ---------------------------------------------------------------------------
+# # Side-effect verification
+# # ---------------------------------------------------------------------------
+
+# def test_ffmpeg_called_once(patched_env):
+#     _call(patched_env["tmp_path"])
+#     patched_env["p_subprocess"].assert_called_once()
+
+
+# def test_temp_file_removed(patched_env):
+#     _call(patched_env["tmp_path"])
+#     patched_env["p_remove"].assert_called_once()
+
+
+# def test_writer_called_per_frame(patched_env):
+#     """cv2.VideoWriter.write() must be called once per video frame."""
+#     n_video_frames = patched_env["video_array"].shape[0]
+#     _call(patched_env["tmp_path"])
+#     assert patched_env["mock_writer"].write.call_count == n_video_frames
+
+
+# # ---------------------------------------------------------------------------
+# # AEM overlay logic
+# # ---------------------------------------------------------------------------
+
+# def test_channel2_increased_by_aem(patched_env):
+#     """Channel index 2 of every written frame must be >= the original (AEM adds energy)."""
+#     original_ch2 = patched_env["video_array"][:, :, :, 2].astype(np.float64)
+
+#     _call(patched_env["tmp_path"])
+
+#     written_frames = [
+#         call_args[0][0]  # first positional arg to write()
+#         for call_args in patched_env["mock_writer"].write.call_args_list
+#     ]
+#     for i, frame in enumerate(written_frames):
+#         assert np.all(frame[:, :, 2].astype(np.float64) >= original_ch2[i] * 0.94), (
+#             f"Frame {i}: channel 2 should not decrease below original × 0.95 brightness factor"
+#         )
+
+
+# # ---------------------------------------------------------------------------
+# # Integration tests — require real files and ffmpeg
+# # Run with: pytest -m integration
+# # ---------------------------------------------------------------------------
+
+# # @pytest.mark.integration
+# # @pytest.mark.parametrize("normalization_mode", [None, "global", "local"])
+# # def test_integration_smoke(tmp_path, normalization_mode):
+# #     """End-to-end run with real audio/video files; verifies output .mp4 is created."""
+# #     out = tmp_path / "out.mp4"
+# #     visualize_aem(
+# #         str(REAL_AUDIO_PATH),
+# #         str(REAL_VIDEO_PATH),
+# #         str(out),
+# #         duration=1.0,
+# #         audio_offset=0.0,
+# #         video_offset=0.0,
+# #         normalization_mode=normalization_mode,
+# #     )
+# #     assert out.exists(), "Output video file was not created"
+# #     assert out.stat().st_size > 0, "Output video file is empty"