xgboost

predict/predict_xgboost.py

@@ -0,0 +1,155 @@
+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
+):
+    # =====================================================
+    # 경로
+    # =====================================================
+
+    features_path = os.path.join(
+        working_dir,
+        FEATURES_FILENAME
+    )
+
+    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 생성
+    # =====================================================
+
+    row = []
+
+    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]
+
+        diff = target.get("difficulty", "unknown")
+
+        print(
+            f"{diff:10} "
+            f"{pred:.1f}"
+        )
+
+
+# =========================================================
+# main
+# =========================================================
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--workingDir",
+        required=True
+    )
+
+    parser.add_argument(
+        "--songno",
+        required=True
+    )
+
+    args = parser.parse_args()
+
+    predict(
+        args.workingDir,
+        args.songno
+    )