File size: 2,475 Bytes
88b5dc0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import itertools
import os
import random
import time
import warnings

import numpy as np

_DEBUG = bool(os.environ.get("DEBUG", False))


class Effect:
    def apply(self, wav: np.ndarray, sr: int):
        """
        Args:
            wav: (T)
            sr: sample rate
        Returns:
            wav: (T) with the same sample rate of `sr`
        """
        raise NotImplementedError

    def __call__(self, wav: np.ndarray, sr: int):
        """
        Args:
            wav: (T)
            sr: sample rate
        Returns:
            wav: (T) with the same sample rate of `sr`
        """
        assert len(wav.shape) == 1, wav.shape

        if _DEBUG:
            start = time.time()
        else:
            start = None

        shape = wav.shape
        assert wav.ndim == 1, f"{self}: Expected wav.ndim == 1, got {wav.ndim}."
        wav = self.apply(wav, sr)
        assert shape == wav.shape, f"{self}: {shape} != {wav.shape}."

        if start is not None:
            end = time.time()
            print(f"{self.__class__.__name__}: {end - start:.3f} sec")

        return wav


class Chain(Effect):
    def __init__(self, *effects):
        super().__init__()

        self.effects = effects

    def apply(self, wav, sr):
        for effect in self.effects:
            wav = effect(wav, sr)
        return wav


class Maybe(Effect):
    def __init__(self, prob, effect):
        super().__init__()

        self.prob = prob
        self.effect = effect

        if _DEBUG:
            warnings.warn("DEBUG mode is on. Maybe -> Must.")
            self.prob = 1

    def apply(self, wav, sr):
        if random.random() > self.prob:
            return wav
        return self.effect(wav, sr)


class Choice(Effect):
    def __init__(self, *effects, **kwargs):
        super().__init__()
        self.effects = effects
        self.kwargs = kwargs

    def apply(self, wav, sr):
        return np.random.choice(self.effects, **self.kwargs)(wav, sr)


class Permutation(Effect):
    def __init__(self, *effects, n: int | None = None):
        super().__init__()
        self.effects = effects
        self.n = n

    def apply(self, wav, sr):
        if self.n is None:
            n = np.random.binomial(len(self.effects), 0.5)
        else:
            n = self.n
        if n == 0:
            return wav
        perms = itertools.permutations(self.effects, n)
        effects = random.choice(list(perms))
        return Chain(*effects)(wav, sr)