164 lines
3.7 KiB
Python
164 lines
3.7 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import joblib
|
|
import numpy as np
|
|
|
|
|
|
# =========================================================
|
|
# 파일명
|
|
# =========================================================
|
|
|
|
FEATURES_FILENAME = "features.json"
|
|
|
|
MODEL_FILENAME = "model.pkl"
|
|
SCALER_FILENAME = "scaler.pkl"
|
|
FEATURE_NAMES_FILENAME = "features.txt"
|
|
|
|
|
|
# =========================================================
|
|
# safe float
|
|
# =========================================================
|
|
|
|
def safe_float(value):
|
|
if value is None:
|
|
return 0.0
|
|
|
|
x = float(value)
|
|
|
|
if not math.isfinite(x):
|
|
return 0.0
|
|
|
|
return x
|
|
|
|
|
|
# =========================================================
|
|
# 예측 함수
|
|
# =========================================================
|
|
|
|
def predict(
|
|
working_dir: str,
|
|
songno: str,
|
|
feature: str = None
|
|
):
|
|
# =====================================================
|
|
# 경로
|
|
# =====================================================
|
|
|
|
features_path = (
|
|
os.path.join(working_dir, FEATURES_FILENAME)
|
|
if feature is None
|
|
else feature
|
|
)
|
|
|
|
model_path = os.path.join(
|
|
working_dir,
|
|
MODEL_FILENAME
|
|
)
|
|
|
|
scaler_path = os.path.join(
|
|
working_dir,
|
|
SCALER_FILENAME
|
|
)
|
|
|
|
feature_names_path = os.path.join(
|
|
working_dir,
|
|
FEATURE_NAMES_FILENAME
|
|
)
|
|
|
|
# =====================================================
|
|
# 모델 로드
|
|
# =====================================================
|
|
|
|
model = joblib.load(model_path)
|
|
scaler = joblib.load(scaler_path)
|
|
|
|
# =====================================================
|
|
# feature 이름 로드
|
|
# =====================================================
|
|
|
|
with open(feature_names_path, "r", encoding="utf-8") as f:
|
|
feature_names = [
|
|
line.strip()
|
|
for line in f.readlines()
|
|
if line.strip()
|
|
]
|
|
|
|
# =====================================================
|
|
# features.json 로드
|
|
# =====================================================
|
|
|
|
with open(features_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
# =====================================================
|
|
# target 찾기
|
|
# =====================================================
|
|
|
|
targets = []
|
|
|
|
for item in data:
|
|
if str(item["songno"]) == str(songno):
|
|
targets.append(item)
|
|
|
|
if len(targets) == 0:
|
|
raise ValueError(f"Chart not found: songno={songno}")
|
|
|
|
# =====================================================
|
|
# feature vector 생성
|
|
# =====================================================
|
|
|
|
results = []
|
|
|
|
for target in targets:
|
|
row = []
|
|
|
|
for k in feature_names:
|
|
value = target.get(k, 0)
|
|
row.append(safe_float(value))
|
|
|
|
X = np.array([row], dtype=np.float32)
|
|
|
|
X = scaler.transform(X)
|
|
|
|
pred = model.predict(X)[0]
|
|
|
|
results.append({
|
|
"songno": str(songno),
|
|
"diff": target.get("difficulty", "unknown"),
|
|
"predicted": round(float(pred), 4)
|
|
})
|
|
|
|
print(json.dumps(results, indent=2, ensure_ascii=False))
|
|
|
|
|
|
# =========================================================
|
|
# main
|
|
# =========================================================
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--workingDir",
|
|
required=True
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--feature",
|
|
required=False
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--songno",
|
|
required=True
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
predict(
|
|
args.workingDir,
|
|
args.songno,
|
|
args.feature
|
|
) |