69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import joblib
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
# =========================================================
|
|
# Configuration (Must match training)
|
|
# =========================================================
|
|
MAX_NOTES = 2000
|
|
FACTOR_COUNT = 4
|
|
INPUT_DIM = MAX_NOTES * FACTOR_COUNT
|
|
|
|
MODEL_FILENAME = "model.pkl"
|
|
SCALER_FILENAME = "scaler.pkl"
|
|
|
|
def safe_float(value):
|
|
try: return float(value)
|
|
except: return 0.0
|
|
|
|
def predict():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--workingDir", required=True)
|
|
parser.add_argument("--songno", required=True)
|
|
parser.add_argument("--factor", required=True)
|
|
args = parser.parse_args()
|
|
|
|
model_path = os.path.join(args.workingDir, MODEL_FILENAME)
|
|
scaler_path = os.path.join(args.workingDir, SCALER_FILENAME)
|
|
|
|
if not os.path.exists(model_path):
|
|
print(f"Model not found: {model_path}")
|
|
return
|
|
|
|
model = joblib.load(model_path)
|
|
scaler = joblib.load(scaler_path)
|
|
|
|
with open(args.factor or (Path(args.workingDir) / 'factors.json'), "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
targets = [item for item in data if str(item["songno"]) == str(args.songno)]
|
|
if not targets:
|
|
print(f"No data found for songno: {args.songno}")
|
|
return
|
|
|
|
results = []
|
|
for item in targets:
|
|
raw_factors = item["factors"]
|
|
vector = np.zeros(INPUT_DIM, dtype=np.float32)
|
|
|
|
for i in range(min(len(raw_factors), MAX_NOTES)):
|
|
for j in range(FACTOR_COUNT):
|
|
vector[i * FACTOR_COUNT + j] = safe_float(raw_factors[i][j])
|
|
|
|
X = scaler.transform([vector])
|
|
pred = model.predict(X)[0]
|
|
|
|
results.append({
|
|
"songno": item["songno"],
|
|
"diff": item["difficulty"],
|
|
"predicted": float(pred)
|
|
})
|
|
|
|
print(json.dumps(results, indent=2))
|
|
|
|
if __name__ == "__main__":
|
|
predict()
|