55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
import os
|
|
from pathlib import Path
|
|
import pandas as pd
|
|
import numpy as np
|
|
import streamlit as st
|
|
|
|
from music_engine.matcher import MusicMatcher
|
|
from music_engine.image_processor import ImageProcessor
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
|
|
@st.cache_resource
|
|
def load_music_engine():
|
|
# Инициализация базы данных и регрессора для музыкального мэтчинга
|
|
db_path = BASE_DIR.parent / "dataset" / "DEAM" / "music_db.csv"
|
|
model_path = BASE_DIR / "music_engine" / "va_regressor.pkl"
|
|
|
|
if not db_path.exists():
|
|
print(f"Музыкальная БД не найдена: {db_path}")
|
|
return None
|
|
|
|
return MusicMatcher(db_path=db_path, model_path=model_path)
|
|
|
|
@st.cache_resource
|
|
def load_image_processor():
|
|
# Модуль обработки визуальных признаков
|
|
model_path = BASE_DIR / "emoset_resnet50_best.pth"
|
|
|
|
# Обработка пути при вызове из корневой директории
|
|
if not model_path.exists():
|
|
model_path = BASE_DIR.parent / "emoset_resnet50_best.pth"
|
|
|
|
return ImageProcessor(model_path=model_path)
|
|
|
|
@st.cache_data
|
|
def load_emoset_data():
|
|
# Выборка данных датасета для вкладки отладки
|
|
dataset_root = BASE_DIR.parent / "dataset" / "EmoSet-118K" / "test"
|
|
|
|
csv_path = dataset_root / "labels.csv"
|
|
img_dir = dataset_root / "images"
|
|
emb_path = BASE_DIR / "emoset_test_embeddings.npy"
|
|
lbl_path = BASE_DIR / "emoset_test_labels.npy"
|
|
|
|
if not all([csv_path.exists(), emb_path.exists(), lbl_path.exists()]):
|
|
print("Тестовые файлы датасета не найдены, вкладка отладки может работать некорректно")
|
|
return None, None, None, None
|
|
|
|
labels_df = pd.read_csv(csv_path)
|
|
|
|
test_filenames = labels_df['filename'].tolist()
|
|
test_embeddings = np.load(emb_path)
|
|
test_labels = np.load(lbl_path)
|
|
|
|
return test_filenames, test_embeddings, test_labels, img_dir |