File size: 1,718 Bytes
71bca69
 
 
c6b6849
71bca69
 
 
a04cce4
 
71bca69
 
 
 
 
 
 
 
 
 
 
580abb7
71bca69
 
 
580abb7
71bca69
 
 
 
 
 
fc0de2c
71bca69
 
 
a04cce4
71bca69
 
 
 
 
 
fc0de2c
71bca69
 
 
fc0de2c
 
 
 
 
71bca69
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import torch
import torchvision.transforms as transforms
from huggingface_hub import snapshot_download
from PIL import Image

MODEL_DIR = snapshot_download(
    "ccmusic-database/chest_falsetto",
    cache_dir="./__pycache__",
)


def toCUDA(x):
    if hasattr(x, "cuda"):
        if torch.cuda.is_available():
            return x.cuda()

    return x


def find_files(folder_path=f"{MODEL_DIR}/examples", ext=".wav"):
    wav_files = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(ext):
                file_path = os.path.join(root, file)
                wav_files.append(file_path)

    return wav_files


def get_modelist(model_dir=MODEL_DIR, assign_model=""):
    try:
        entries = os.listdir(model_dir)
    except OSError as e:
        print(f"Cannot access {model_dir}: {e}")
        return

    output = []
    for entry in entries:
        full_path = os.path.join(model_dir, entry)
        if entry == ".git" or entry == "examples":
            print(f"Skip .git / examples dir: {full_path}")
            continue

        if os.path.isdir(full_path):
            model = os.path.basename(full_path)
            if assign_model and assign_model.lower() in model:
                output.insert(0, model)
            else:
                output.append(model)

    return output


def embed_img(img_path: str, input_size=224):
    transform = transforms.Compose(
        [
            transforms.Resize([input_size, input_size]),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )
    img = Image.open(img_path).convert("RGB")
    return transform(img).unsqueeze(0)