71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import joblib
|
|
import numpy as np
|
|
import warnings
|
|
|
|
# 경고 무시 (Feature name 관련 경고 제거)
|
|
warnings.filterwarnings("ignore", category=UserWarning)
|
|
|
|
# =========================================================
|
|
# 파일명
|
|
# =========================================================
|
|
|
|
FEATURES_FILENAME = "features.json"
|
|
MODEL_FILENAME = "model.pkl"
|
|
SCALER_FILENAME = "scaler.pkl"
|
|
FEATURE_NAMES_FILENAME = "features.txt"
|
|
|
|
def safe_float(value):
|
|
if value is None: return 0.0
|
|
x = float(value)
|
|
return x if math.isfinite(x) else 0.0
|
|
|
|
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)
|
|
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model not found at {model_path}")
|
|
|
|
model = joblib.load(model_path)
|
|
scaler = joblib.load(scaler_path)
|
|
|
|
with open(feature_names_path, "r", encoding="utf-8") as f:
|
|
feature_names = [line.strip() for line in f.readlines() if line.strip()]
|
|
|
|
with open(features_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
targets = [item for item in data if str(item["songno"]) == str(songno)]
|
|
if len(targets) == 0:
|
|
raise ValueError(f"Chart not found: songno={songno}")
|
|
|
|
results = []
|
|
for target in targets:
|
|
row = [safe_float(target.get(k, 0)) for k in feature_names]
|
|
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))
|
|
|
|
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)
|