sma_processing/extractor.py at master · aligokce/sma_processing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
from pathlib import Path

import numpy as np

from dataset import SMIRDataset
from dataset.utils import emulate_scene
from features import rent, shd, srf
from microphone import microphones

from omegaconf import OmegaConf, open_dict


class Extractor:
    def __init__(self, params: OmegaConf, dataset: SMIRDataset) -> None:
        """
        Feature extractor via emulation featuring anechoic sounds and
        multichannel impulse response dataset.

        Parameters
        ----------
        params : OmegaConf
            Parameters for feature extraction
        dataset : SMIRDataset
            Spherical microphone impulse response dataset
        """
        self.params = self._load_params(params)
        self.dataset = dataset

    def job(self, anechoic_path, scenario, save=True):
        """
        RENT extraction job

        Parameters
        ----------
        anechoic_path : Union[str, bytes, PathLike]
            Path for anechoic sound file path
        scenario : Dict
            Dictionary describing the current emulation scenario
        save : bool, optional
            Results are saved when true, returned when false, by default True
        output_shd : bool, optional
            Whether to additionaly output intermediate SHD matrix, by default False

        Returns
        -------
        None, NDArray or (NDArray, NDArray)
            Nothing is returned if `save` is enabled.
            RENT matrix is returned if `save` is disabled.
            RENT and SHD matrices are returned if
            `output_shd` is enabled and `save` is disabled.
        """
        y = self.load_signal(anechoic_path, scenario)

        res = rent_pipe(
            y,
            **self.params,
            mic=microphones[scenario['mic']]
        )

        if save:
            return self.save_result(res, anechoic_path, scenario)
        return res

    def load_signal(self, sig_path, scenario):
        # NOTE: Should this function belong to the SMIRDataset class?
        # Calculate gain
        mic_pos = scenario['mic_pos']
        src_pos = scenario['src_pos']
        gain = self.dataset.calculate_gain(src_pos, mic_pos)

        # Generate IR paths
        ir_paths = self.dataset.generate_ir_paths(**scenario)

        sig = emulate_scene(sig_path, gain, ir_paths)
        return sig

    def save_result(self, result, anechoic_path, scenario):
        save_folder = ""
        self._save_npy(result, save_folder, anechoic_path, **scenario)

    @staticmethod
    def _load_params(params: OmegaConf):
        p = params.copy()
        with open_dict(p):
            p['fimin'] = int(
                round(params['fl'] / params['fs'] * params['n_fft']))
            p['fimax'] = int(
                round(params['fh'] / params['fs'] * params['n_fft']))
        return p

    @staticmethod
    def _save_npy(npy, save_folder, anechoic_path, mic_pos, src_pos, src_dir, **kwargs):
        # TODO: How to split all variables from here while anechoic filename contains "_"
        # Example output filename: OA-09_PA-19_W__mahler_vl1b_6
        save_name = f"{mic_pos}_{src_pos}_{src_dir}__{anechoic_path.stem}.npy"
        save_path = Path.cwd() / save_folder / save_name

        save_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(save_path, npy)


def rent_pipe(
    y,
    fs,     # shd
    n_fft,
    olap,
    n_shd,
    fimin,
    fimax,
    j_nu,
    mic,
    n_pix,  # srf
    **kwargs
):
    Anm = shd.extract(y, fs, n_fft, olap, n_shd, fimin, fimax, j_nu, mic)
    S = srf.extract(Anm, n_pix)
    R = rent.extract(S, n_shd, fimin, fimax, j_nu)

    return R