ccmusic-database/CTIS

1import os2import torch3import torchvision.transforms as transforms4from PIL import Image5 6EN_US = os.getenv("LANG") != "zh_CN.UTF-8"7 8ZH2EN = {9    "上传录音": "Upload a recording",10    "选择模型": "Select a model",11    "状态栏": "Status",12    "音频文件名": "Audio filename",13    "中国乐器识别": "Chinese instrument recognition",14    "建议录音时长保持在 3s 左右": "It is recommended to keep the recording length around 3s.",15    "引用": "Cite",16}17 18if EN_US:19    import huggingface_hub20 21    MODEL_DIR = huggingface_hub.snapshot_download(22        "ccmusic-database/CTIS",23        cache_dir="./__pycache__",24    )25 26else:27    import modelscope28 29    MODEL_DIR = modelscope.snapshot_download(30        "ccmusic-database/CTIS",31        cache_dir="./__pycache__",32    )33 34 35def _L(zh_txt: str):36    return ZH2EN[zh_txt] if EN_US else zh_txt37 38 39def toCUDA(x):40    if hasattr(x, "cuda"):41        if torch.cuda.is_available():42            return x.cuda()43 44    return x45 46 47def find_files(folder_path=f"{MODEL_DIR}/examples", ext=".wav"):48    wav_files = []49    for root, _, files in os.walk(folder_path):50        for file in files:51            if file.endswith(ext):52                file_path = os.path.join(root, file)53                wav_files.append(file_path)54 55    return wav_files56 57 58def get_modelist(model_dir=MODEL_DIR, assign_model=""):59    output = []60    for entry in os.listdir(model_dir):61        # 获取完整路径62        full_path = os.path.join(model_dir, entry)63        # 跳过'.git'文件夹64        if entry == ".git" or entry == "examples":65            print(f"跳过 .git 或 examples 文件夹: {full_path}")66            continue67 68        # 检查条目是文件还是目录69        if os.path.isdir(full_path):70            model = os.path.basename(full_path)71            if assign_model and assign_model.lower() in model:72                output.insert(0, model)73            else:74                output.append(model)75 76    return output77 78 79def embed_img(img_path: str, input_size=224):80    transform = transforms.Compose(81        [82            transforms.Resize([input_size, input_size]),83            transforms.ToTensor(),84            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),85        ]86    )87    img = Image.open(img_path).convert("RGB")88    return transform(img).unsqueeze(0)89