image

predict/factor/predict_lightgbm.py

@@ -0,0 +1,68 @@
+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()