|
import torch |
|
from model.bart import BartCaptionModel |
|
from utils.audio_utils import load_audio, STR_CH_FIRST |
|
from typing import Dict, List, Any |
|
import numpy as np |
|
import librosa |
|
import os |
|
import json |
|
def preprocess_audio(audio_signal, sr, duration=10, target_sr=16000): |
|
n_samples = int(duration * target_sr) |
|
audio = librosa.to_mono(audio_signal) |
|
audio = librosa.resample(audio, orig_sr = sr, target_sr = target_sr) |
|
|
|
if len(audio.shape) == 2: |
|
audio = audio.mean(0, False) |
|
input_size = int(n_samples) |
|
if audio.shape[-1] < input_size: |
|
pad = np.zeros(input_size) |
|
pad[: audio.shape[-1]] = audio |
|
audio = pad |
|
ceil = int(audio.shape[-1] // n_samples) |
|
audio = torch.from_numpy(np.stack(np.split(audio[:ceil * n_samples], ceil)).astype('float32')) |
|
return audio |
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
if os.path.isfile("transfer.pth") == False: |
|
torch.hub.download_url_to_file('https://huggingface.co/seungheondoh/lp-music-caps/resolve/main/transfer.pth', 'transfer.pth') |
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu" |
|
self.model = BartCaptionModel(max_length = 128) |
|
pretrained_object = torch.load('./transfer.pth', map_location='cpu') |
|
state_dict = pretrained_object['state_dict'] |
|
self.model.load_state_dict(state_dict) |
|
if torch.cuda.is_available(): |
|
torch.cuda.set_device(self.device) |
|
self.model = self.model.cuda(self.device) |
|
|
|
def _captioning(self, audio_tensor): |
|
if self.device is not None: |
|
audio_tensor = audio_tensor.to(self.device) |
|
|
|
with torch.no_grad(): |
|
output = self.model.generate( |
|
samples=audio_tensor, |
|
num_beams=5, |
|
) |
|
inference = "" |
|
number_of_chunks = range(audio_tensor.shape[0]) |
|
for chunk, text in zip(number_of_chunks, output): |
|
time = f"[{chunk * 10}:00-{(chunk + 1) * 10}:00]" |
|
inference += f"{time}\n{text} \n \n" |
|
return inference |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
data = json.loads(data["inputs"]) |
|
array = np.array(data['audio_list'], dtype=data["audio_dtype"]) |
|
array_shape = data['audio_shape'] |
|
input_audio = array.reshape(array_shape) |
|
sr = data["sampling_rate"] |
|
|
|
preprocessed_audio = preprocess_audio(input_audio, sr) |
|
|
|
return self._captioning(preprocessed_audio) |
|
""" |
|
if __name__ == "__main__": |
|
import numpy as np |
|
from scipy.io.wavfile import write as wav_write |
|
from huggingface_hub import InferenceApi |
|
|
|
handler = EndpointHandler() |
|
audio_path = "folk.wav" |
|
np_audio, sr = librosa.load(audio_path, sr=44100) |
|
|
|
np_list = np_audio.tolist() |
|
np_shape = np_audio.shape |
|
np_dtype = np_audio.dtype.name |
|
|
|
request = json.dumps({ |
|
"audio_list": np_list, |
|
"audio_shape": np_shape, |
|
"audio_dtype": np_dtype, |
|
"sampling_rate": sr |
|
}) |
|
|
|
print(f"Loaded {audio_path} with sample rate {sr}") |
|
print(handler.__call__({"payload": request})) |
|
""" |
