vocal_isolation/loader.py

# Standard Library
import os

# Explicit Typing
from typing import Tuple
from numpy import ndarray

# Third-party
import librosa
import torch


class Loader:
    """Loading sound files into a usable format for pytorch"""

    def __init__(self, INPUT_FOLDER, OUTPUT_FOLDER):
        self.input = INPUT_FOLDER
        self.output = OUTPUT_FOLDER

    def load_wav(self, name) -> Tuple[torch.Tensor, int]:
        music_array, samplerate = librosa.load(
            os.path.join(self.input, name + ".wav"), mono=False, sr=44100
        )
        music_tensor = torch.tensor(music_array, dtype=torch.float32)
        return music_tensor, samplerate

    def prepare_uploaded_file(self, uploaded_file) -> Tuple[torch.Tensor, int]:
        music_array, samplerate = librosa.load(uploaded_file, mono=False, sr=44100)

        music_tensor = torch.tensor(music_array, dtype=torch.float32)

        return music_tensor, samplerate