Audio Classification
autrainer
audio
ecoacoustic-tagging
HearTheSpecies
ecoacoustics
File size: 4,187 Bytes
5374684
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
120
121
122
123
124
125
126
127
128
129
130
import warnings

import torch
from transformers import ClapAudioModel, ClapAudioModelWithProjection, ClapFeatureExtractor, ClapProcessor

from autrainer.models.abstract_model import AbstractModel
from autrainer.models.ffnn import FFNN


class CLAPBackbone(AbstractModel):
    def __init__(
        self,
        model_name,
        freeze_extractor: bool = True,
        time_pooling: bool = True,
    ) -> None:
        self.model_name = model_name
        self.freeze_extractor = freeze_extractor
        self.time_pooling = time_pooling
        
        model = ClapAudioModelWithProjection.from_pretrained(self.model_name)
        super().__init__(output_dim=model.config.hidden_size)

        self.model = model.audio_model.audio_encoder
        # self.model = model
        # print(self.model)

        if self.freeze_extractor:
            for param in self.model.parameters():
                param.requires_grad = False

    def embeddings(self, x: torch.Tensor) -> torch.Tensor:
        inputs = x
        is_longer = torch.tensor([False])

        x = self.model(input_features=inputs, is_longer=is_longer).last_hidden_state

        # Flatten and transpose for the embeddings
        x = x.flatten(2).transpose(1, 2)

        if self.time_pooling:
            x = x.mean(1)

        return x

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.embeddings(features)


class CLAPFFNN(AbstractModel):
    def __init__(
        self,
        output_dim: int,
        model_name: str,
        freeze_extractor: bool,
        hidden_size: int,
        num_layers: int = 2,
        dropout: float = 0.5,
    ) -> None:
        """CLAP model with FFNN frontend adapted for audio classification.
        For more information, see: https://huggingface.co/docs/transformers/en/model_doc/clap#clap

        Args:
            output_dim: Output dimension of the FFNN.
            model_name: Name of the model loaded from Huggingface.
            freeze_extractor: Whether to freeze the feature extractor.
            hidden_size: Hidden size of the FFNN.
            num_layers: Number of layers of the FFNN. Defaults to 2.
            dropout: Dropout rate. Defaults to 0.5.
        """
        super().__init__(output_dim)
        self.model_name = model_name
        self.freeze_extractor = freeze_extractor
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.backbone = CLAPBackbone(
            model_name=model_name,
            freeze_extractor=freeze_extractor,
            time_pooling=True,
        )
        self.frontend = FFNN(
            input_size=self.backbone.output_dim,
            hidden_size=hidden_size,
            output_dim=output_dim,
            num_layers=num_layers,
            dropout=dropout,
        )

    def embeddings(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.frontend(self.embeddings(features)) 


if __name__=='__main__':
    output_dim = 4
    model_name = "laion/clap-htsat-fused"
    freeze_extractor = True
    time_pooling = True
    hidden_size = 512

    model = CLAPFFNN(
        output_dim=output_dim,
        model_name = model_name,
        freeze_extractor = freeze_extractor,
        hidden_size=hidden_size
    )

    feature_extractor = ClapFeatureExtractor.from_pretrained('laion/clap-htsat-unfused')
    # processor = ClapProcessor.from_pretrained('laion/clap-htsat-unfused')

    import librosa
    a, sr = librosa.load("/path/to/example.wav", sr=48000)
    print(a.shape, sr)
    audio = torch.tensor(a)

    # inputs = processor(audios=audio, sampling_rate=48000, return_tensors="pt")
    # print("Inputs:", inputs['input_features'].shape)
    extracted = feature_extractor(audio, sampling_rate=48000, return_tensors='pt')
    print("Extracted: ", extracted['input_features'].shape)
    extracted = extracted['input_features']
    # print(type(inputs))
    print(type(extracted))

    # features = extracted[list(extracted.keys())[0]][0].unsqueeze(0)
    out = model(extracted)
    print(out)
    print(out.shape)