Upload folder using huggingface_hub

README.md

@@ -1,12 +1,12 @@
 ---
-title: Kws Fp Test
-emoji: ⚡
-colorFrom: yellow
-colorTo: gray
+title: KWS FP Test
+emoji: 🎤
+colorFrom: green
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.7.0
+sdk_version: 6.5.1
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+KWS FP Test Space

app.py

@@ -0,0 +1,6 @@
+from gradio_kws_test import build_demo
+
+demo = build_demo()
+
+if __name__ == "__main__":
+    demo.launch()

checkpoint_io.py

@@ -0,0 +1,119 @@
+# checkpoint_io.py: KWS 체크포인트 Safetensors 저장/로드 + 기존 .pth 호환.
+# Single Source of Truth: 학습은 .safetensors + training_config.json, 로드는 .safetensors 우선 후 .pth 폴백.
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+
+def resolve_checkpoint_path(path: str) -> Tuple[str, str]:
+    """
+    체크포인트 경로에서 (state 파일 경로, config 파일 경로) 반환.
+    .safetensors 우선, 없으면 .pth. config는 항상 같은 디렉터리의 training_config.json.
+    """
+    path = os.path.normpath(path)
+    if os.path.isdir(path):
+        base_dir = path
+        base_name = "best_kws_model"
+    else:
+        base_dir = os.path.dirname(path) or "."
+        base_name = os.path.splitext(os.path.basename(path))[0] or "best_kws_model"
+    config_path = os.path.join(base_dir, "training_config.json")
+    for ext in (".safetensors", ".pth"):
+        candidate = os.path.join(base_dir, base_name + ext)
+        if os.path.isfile(candidate):
+            return candidate, config_path
+    return os.path.join(base_dir, base_name + ".pth"), config_path
+
+
+def load_state_dict(path: str, map_location: Optional[Any] = None) -> Dict[str, torch.Tensor]:
+    """
+    path(또는 동일 디렉터리 내 best_kws_model.*)에서 state_dict만 로드.
+    .safetensors 우선, 없으면 .pth (전체 체크포인트면 'model' 키 사용).
+    """
+    state_path, _ = resolve_checkpoint_path(path)
+    if not os.path.isfile(state_path):
+        raise FileNotFoundError(f"체크포인트 없음: {state_path}")
+    if state_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        device: str
+        if map_location is None:
+            device = "cpu"
+        elif isinstance(map_location, torch.device):
+            if map_location.type == "cuda":
+                device = f"cuda:{map_location.index}" if map_location.index is not None else "cuda"
+            else:
+                device = map_location.type
+        elif isinstance(map_location, str):
+            device = map_location
+        else:
+            device = "cpu"
+        return load_file(state_path, device=device)
+    ckpt = torch.load(state_path, map_location=map_location)
+    if isinstance(ckpt, dict) and "model" in ckpt:
+        return ckpt["model"]
+    return ckpt
+
+
+def load_label_map(path: str) -> Optional[Dict[str, int]]:
+    """같은 디렉터리의 training_config.json에서 label_map 반환. 없으면 None."""
+    _, config_path = resolve_checkpoint_path(path)
+    if not os.path.isfile(config_path):
+        return None
+    try:
+        with open(config_path, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+        return cfg.get("label_map")
+    except Exception:
+        return None
+
+
+def load_checkpoint(
+    path: str,
+    map_location: Optional[Any] = None,
+    default_label_map: Optional[Dict[str, int]] = None,
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, int], List[str]]:
+    """
+    state_dict, label_map, class_names 반환.
+    label_map은 training_config.json 우선, 없으면 .pth 내부, 없으면 default_label_map.
+    """
+    state = load_state_dict(path, map_location=map_location)
+    label_map = load_label_map(path)
+    if label_map is None:
+        state_path, _ = resolve_checkpoint_path(path)
+        if state_path.endswith(".pth") and os.path.isfile(state_path):
+            ckpt = torch.load(state_path, map_location=map_location)
+            if isinstance(ckpt, dict) and "label_map" in ckpt:
+                label_map = ckpt["label_map"]
+        if label_map is None:
+            label_map = default_label_map or {"normal": 0, "help_me": 1, "save_me": 2}
+    num_classes = len(label_map)
+    id_to_name = {v: k for k, v in label_map.items()}
+    class_names: List[str] = [id_to_name[i] for i in range(num_classes)]
+    return state, label_map, class_names
+
+
+def save_checkpoint(
+    output_dir: str,
+    state_dict: Dict[str, torch.Tensor],
+    config_dict: Dict[str, Any],
+    *,
+    also_save_pth: bool = False,
+) -> None:
+    """
+    state_dict를 best_kws_model.safetensors로, config_dict를 training_config.json으로 저장.
+    also_save_pth=True면 기존 호환용 best_kws_model.pth도 저장 (state_dict만).
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    sf_path = os.path.join(output_dir, "best_kws_model.safetensors")
+    from safetensors.torch import save_file
+    save_file(state_dict, sf_path)
+    config_path = os.path.join(output_dir, "training_config.json")
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(config_dict, f, indent=2, ensure_ascii=False)
+    if also_save_pth:
+        pth_path = os.path.join(output_dir, "best_kws_model.pth")
+        torch.save(state_dict, pth_path)

gradio_kws_test.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+import librosa
+import numpy as np
+
+from kws_inference import KWSLongInference, SR_MODEL
+from train_kws import WINDOW_SEC
+
+
+def _default_model_path() -> str:
+    candidates = [
+        Path("./kws_models_fpfix/best_kws_model.safetensors"),
+        Path("./kws_models/best_kws_model.safetensors"),
+        Path("./kws_models_fixed/best_kws_model.safetensors"),
+        Path("./kws_models/best_kws_model.pth"),
+    ]
+    for p in candidates:
+        if p.exists():
+            return str(p)
+    return str(candidates[0])
+
+
+_ENGINE_CACHE: dict[tuple[str, float, float], KWSLongInference] = {}
+
+
+def _get_engine(model_path: str, step_sec: float, rms_threshold: float) -> KWSLongInference:
+    key = (str(Path(model_path).resolve()), float(step_sec), float(rms_threshold))
+    if key not in _ENGINE_CACHE:
+        _ENGINE_CACHE[key] = KWSLongInference(
+            checkpoint_path=model_path,
+            step_sec=step_sec,
+            rms_threshold=rms_threshold,
+        )
+    return _ENGINE_CACHE[key]
+
+
+def _simulate_alerts(
+    window_results: list[Any],
+    threshold: float,
+    n_consecutive: int,
+    cooldown_sec: float,
+) -> list[dict[str, Any]]:
+    alerts: list[dict[str, Any]] = []
+    consecutive = 0
+    last_alert_sec = -1e9
+    for w in window_results:
+        help_p = float(w.probs.get("help_me", 0.0))
+        save_p = float(w.probs.get("save_me", 0.0))
+        kw_p = max(help_p, save_p)
+        kw_label = "save_me" if save_p >= help_p else "help_me"
+        in_cooldown = (w.start_sec - last_alert_sec) < cooldown_sec
+        if in_cooldown:
+            consecutive = 0
+            continue
+        if kw_p >= threshold:
+            consecutive += 1
+        else:
+            consecutive = 0
+        if consecutive >= n_consecutive:
+            alerts.append(
+                {
+                    "t_sec": round(float(w.start_sec), 2),
+                    "label": kw_label,
+                    "score": round(float(kw_p), 4),
+                    "rule": f"{n_consecutive}x >= {threshold:.2f}",
+                }
+            )
+            last_alert_sec = float(w.start_sec)
+            consecutive = 0
+    return alerts
+
+
+def run_test(
+    audio_path: str,
+    model_path: str,
+    step_sec: float,
+    rms_threshold: float,
+    alert_threshold: float,
+    n_consecutive: int,
+    cooldown_sec: float,
+):
+    if not audio_path:
+        raise gr.Error("오디오 파일(또는 마이크 입력)을 넣어주세요.")
+    if not model_path or not Path(model_path).exists():
+        raise gr.Error(f"모델 파일이 없습니다: {model_path}")
+
+    engine = _get_engine(model_path=model_path, step_sec=step_sec, rms_threshold=rms_threshold)
+    result = engine.predict_long(audio_path)
+
+    rows: list[list[Any]] = []
+    for i, w in enumerate(result.window_results, start=1):
+        rows.append(
+            [
+                i,
+                round(float(w.start_sec), 3),
+                round(float(w.end_sec), 3),
+                w.class_name,
+                round(float(w.confidence), 4),
+                round(float(w.probs.get("normal", 0.0)), 4),
+                round(float(w.probs.get("help_me", 0.0)), 4),
+                round(float(w.probs.get("save_me", 0.0)), 4),
+            ]
+        )
+
+    alerts = _simulate_alerts(
+        window_results=result.window_results,
+        threshold=alert_threshold,
+        n_consecutive=n_consecutive,
+        cooldown_sec=cooldown_sec,
+    )
+
+    agg = result.aggregated_probs
+    summary = (
+        f"### 추론 요약\n"
+        f"- Aggregated label: `{result.aggregated_label}`\n"
+        f"- Duration: `{result.duration_sec:.2f}s`\n"
+        f"- Aggregated probs: normal `{agg.get('normal', 0.0):.4f}` | "
+        f"help_me `{agg.get('help_me', 0.0):.4f}` | "
+        f"save_me `{agg.get('save_me', 0.0):.4f}`\n"
+        f"- Windows: `{len(result.window_results)}` (step `{step_sec:.2f}s`, rms_threshold `{rms_threshold:.4f}`)"
+    )
+
+    if alerts:
+        alert_lines = [
+            f"- t={a['t_sec']:.2f}s | `{a['label']}` | score={a['score']:.4f} | {a['rule']}"
+            for a in alerts
+        ]
+        alerts_md = "### 경보 시뮬레이션\n" + "\n".join(alert_lines)
+    else:
+        alerts_md = "### 경보 시뮬레이션\n- 조건을 만족한 경보 없음"
+
+    details = {
+        "aggregated_label": result.aggregated_label,
+        "aggregated_probs": result.aggregated_probs,
+        "duration_sec": result.duration_sec,
+        "num_windows": len(result.window_results),
+        "alerts": alerts,
+    }
+    return summary, alerts_md, rows, details
+
+
+def _init_stream_state() -> dict[str, Any]:
+    return {
+        "buffer": np.array([], dtype=np.float32),
+        "consecutive": 0,
+        "last_alert_time": -1e9,
+        "events": [],
+        "last_probs": {"normal": 1.0, "help_me": 0.0, "save_me": 0.0},
+        "last_decision": "NORMAL",
+    }
+
+
+def _normalize_chunk(chunk: Any) -> tuple[int, np.ndarray] | None:
+    if chunk is None:
+        return None
+    if isinstance(chunk, np.ndarray):
+        y = np.asarray(chunk, dtype=np.float32)
+        if y.ndim == 2:
+            y = y.mean(axis=1 if y.shape[1] <= 8 else 0)
+        return SR_MODEL, y.astype(np.float32)
+    if isinstance(chunk, (tuple, list)) and len(chunk) == 2:
+        sr, y = chunk
+        y = np.asarray(y, dtype=np.float32)
+        if y.ndim == 2:
+            y = y.mean(axis=1 if y.shape[1] <= 8 else 0)
+        if y.dtype == np.int16:
+            y = y.astype(np.float32) / 32768.0
+        elif y.dtype == np.int32:
+            y = y.astype(np.float32) / 2147483648.0
+        return int(sr), y.astype(np.float32)
+    if isinstance(chunk, dict):
+        # gradio 버전에 따라 {"sample_rate": ..., "data": ...} 혹은 {"path": ...} 형태가 올 수 있음
+        if "sample_rate" in chunk and ("data" in chunk or "array" in chunk):
+            sr = int(chunk.get("sample_rate", SR_MODEL))
+            y = np.asarray(chunk.get("data", chunk.get("array")), dtype=np.float32)
+            if y.ndim == 2:
+                y = y.mean(axis=1 if y.shape[1] <= 8 else 0)
+            return sr, y.astype(np.float32)
+        if "path" in chunk and chunk["path"]:
+            try:
+                y, sr = librosa.load(str(chunk["path"]), sr=None, mono=True)
+                return int(sr), y.astype(np.float32)
+            except Exception:
+                return None
+    return None
+
+
+def stream_infer(
+    state: dict[str, Any] | None,
+    chunk: Any,
+    model_path: str,
+    step_sec: float,
+    rms_threshold: float,
+    alert_threshold: float,
+    margin_threshold: float,
+    n_consecutive: int,
+    cooldown_sec: float,
+):
+    st = state if isinstance(state, dict) else _init_stream_state()
+    parsed = _normalize_chunk(chunk)
+    if parsed is None:
+        probs = st.get("last_probs", {"normal": 1.0, "help_me": 0.0, "save_me": 0.0})
+        return st, "대기 중...", st.get("last_decision", "NORMAL"), probs, list(st.get("events", []))
+
+    sr, y = parsed
+    if len(y) == 0:
+        probs = st.get("last_probs", {"normal": 1.0, "help_me": 0.0, "save_me": 0.0})
+        return st, "입력 없음", st.get("last_decision", "NORMAL"), probs, list(st.get("events", []))
+
+    if sr != SR_MODEL:
+        y = librosa.resample(y, orig_sr=sr, target_sr=SR_MODEL)
+
+    buffer = np.concatenate((st["buffer"], y))
+    max_len = int(SR_MODEL * WINDOW_SEC)
+    if len(buffer) > max_len:
+        buffer = buffer[-max_len:]
+    st["buffer"] = buffer
+
+    if len(buffer) < int(SR_MODEL * 0.6):
+        probs = st.get("last_probs", {"normal": 1.0, "help_me": 0.0, "save_me": 0.0})
+        return st, "데이터 수집 중...", st.get("last_decision", "NORMAL"), probs, list(st.get("events", []))
+
+    engine = _get_engine(model_path=model_path, step_sec=step_sec, rms_threshold=rms_threshold)
+    result = engine.predict_long((SR_MODEL, buffer))
+    probs = {
+        "normal": float(result.aggregated_probs.get("normal", 0.0)),
+        "help_me": float(result.aggregated_probs.get("help_me", 0.0)),
+        "save_me": float(result.aggregated_probs.get("save_me", 0.0)),
+    }
+    st["last_probs"] = probs
+
+    now = time.time()
+    help_p, save_p = probs["help_me"], probs["save_me"]
+    normal_p = probs["normal"]
+    kw_score = max(help_p, save_p)
+    kw_label = "save_me" if save_p >= help_p else "help_me"
+    margin = kw_score - normal_p
+    in_cooldown = (now - float(st["last_alert_time"])) < cooldown_sec
+
+    meets_score = kw_score >= alert_threshold
+    meets_margin = margin >= margin_threshold
+
+    if not in_cooldown and meets_score and meets_margin:
+        st["consecutive"] = int(st["consecutive"]) + 1
+    else:
+        st["consecutive"] = 0
+
+    status = (
+        f"normal={probs['normal']:.3f} | help_me={probs['help_me']:.3f} | "
+        f"save_me={probs['save_me']:.3f} | margin={margin:.3f} | "
+        f"consec={st['consecutive']}/{n_consecutive}"
+    )
+    if in_cooldown:
+        remain = max(0.0, cooldown_sec - (now - float(st["last_alert_time"])))
+        status = f"쿨다운 {remain:.1f}s | " + status
+
+    decision = "NORMAL"
+    probs_txt = f"n={normal_p:.3f} h={help_p:.3f} s={save_p:.3f} m={margin:.3f}"
+
+    if not in_cooldown and meets_score and not meets_margin:
+        decision = "HOLD"
+        event = (
+            f"{time.strftime('%H:%M:%S')} HOLD {kw_label} "
+            f"score={kw_score:.3f} {probs_txt}"
+        )
+        st["events"] = ([event] + list(st.get("events", [])))[:30]
+    elif in_cooldown:
+        decision = "COOLDOWN"
+    elif meets_score and meets_margin:
+        decision = "CANDIDATE"
+
+    if not in_cooldown and int(st["consecutive"]) >= int(n_consecutive):
+        st["last_alert_time"] = now
+        st["consecutive"] = 0
+        event = f"{time.strftime('%H:%M:%S')} ALERT {kw_label} score={kw_score:.3f} {probs_txt}"
+        st["events"] = ([event] + list(st.get("events", [])))[:30]
+        status = "ALERT 발생 | " + status
+        decision = "ALERT"
+
+    st["last_decision"] = decision
+    return st, status, decision, probs, list(st.get("events", []))
+
+
+def build_demo() -> gr.Blocks:
+    with gr.Blocks(title="KWS FP Test Page") as demo:
+        gr.Markdown("# KWS FP Test Page")
+        gr.Markdown("파일 테스트 + 브라우저 마이크 실시간 분석")
+
+        model_path = gr.Textbox(label="모델 경로", value=_default_model_path())
+
+        with gr.Row():
+            step_sec = gr.Slider(label="슬라이딩 스텝(초)", minimum=0.1, maximum=1.0, value=0.25, step=0.05)
+            rms_threshold = gr.Slider(label="RMS 임계값", minimum=0.0, maximum=0.05, value=0.005, step=0.001)
+            alert_threshold = gr.Slider(label="경보 임계값", minimum=0.5, maximum=0.99, value=0.9, step=0.01)
+            margin_threshold = gr.Slider(label="마진 임계값(키워드-normal)", minimum=0.0, maximum=0.6, value=0.2, step=0.01)
+            n_consecutive = gr.Slider(label="연속 횟수", minimum=1, maximum=5, value=3, step=1)
+            cooldown_sec = gr.Slider(label="쿨다운(초)", minimum=0, maximum=10, value=5, step=0.5)
+
+        with gr.Tab("파일 테스트"):
+            audio = gr.Audio(label="테스트 오디오", type="filepath", sources=["upload", "microphone"])
+            run_btn = gr.Button("테스트 실행", variant="primary")
+            summary = gr.Markdown()
+            alerts_md = gr.Markdown()
+            table = gr.Dataframe(
+                headers=["idx", "start_sec", "end_sec", "pred", "conf", "normal", "help_me", "save_me"],
+                datatype=["number", "number", "number", "str", "number", "number", "number", "number"],
+                label="윈도우별 결과",
+            )
+            details = gr.JSON(label="상세 JSON")
+            run_btn.click(
+                fn=run_test,
+                inputs=[audio, model_path, step_sec, rms_threshold, alert_threshold, n_consecutive, cooldown_sec],
+                outputs=[summary, alerts_md, table, details],
+            )
+
+        with gr.Tab("실시간 마이크"):
+            gr.Markdown("마이크를 켠 상태에서 실시간으로 확률/경보 이력을 봅니다.")
+            st = gr.State(value=_init_stream_state())
+            live_audio = gr.Audio(label="마이크 실시간 입력", type="numpy", sources=["microphone"], streaming=True)
+            live_status = gr.Textbox(label="실시간 상태")
+            live_decision = gr.Textbox(label="최종 판단")
+            live_probs = gr.JSON(label="현재 확률")
+            live_events = gr.JSON(label="경보 이력(최신순)")
+            live_audio.stream(
+                fn=stream_infer,
+                inputs=[
+                    st,
+                    live_audio,
+                    model_path,
+                    step_sec,
+                    rms_threshold,
+                    alert_threshold,
+                    margin_threshold,
+                    n_consecutive,
+                    cooldown_sec,
+                ],
+                outputs=[st, live_status, live_decision, live_probs, live_events],
+                show_progress=False,
+                time_limit=300,
+            )
+    return demo
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="KWS FP 테스트용 Gradio 페이지")
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=7861)
+    parser.add_argument("--share", action="store_true")
+    args = parser.parse_args()
+
+    demo = build_demo()
+    demo.launch(server_name=args.host, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()

kws_inference.py

@@ -0,0 +1,197 @@
+"""
+KWS 긴 오디오 추론: 슬라이딩 윈도우 + Max-over-windows 집계.
+[업데이트] 소리 크기(RMS) 필터링 추가로 무음/노이즈 구간 오탐 방지.
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import librosa
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+# train_kws와 동일 파라미터로 mel·모델 호환 유지
+from train_kws import (
+    HOP_LENGTH,
+    KWSModel,
+    MAX_TIME_FRAMES,
+    N_FFT,
+    N_MELS,
+    SR_MODEL,
+    WINDOW_SEC,
+)
+
+# 학습 시 KWSDataset._preprocess와 동일 (일반음성→위급 오탐 방지)
+RMS_NORM_TARGET = 0.05
+PRE_EMPHASIS = 0.97
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+@dataclass
+class WindowResult:
+    """한 윈도우(구간) 추론 결과."""
+    start_sec: float
+    end_sec: float
+    pred_id: int
+    class_name: str
+    probs: Dict[str, float]
+    confidence: float
+
+@dataclass
+class LongAudioResult:
+    """긴 오디오 전체 추론 결과."""
+    aggregated_label: str
+    aggregated_probs: Dict[str, float]
+    aggregated_pred_id: int
+    window_results: List[WindowResult]
+    duration_sec: float
+
+def _load_checkpoint(path: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, int], List[str]]:
+    from checkpoint_io import load_checkpoint
+    state, label_map, class_names = load_checkpoint(path, map_location=DEVICE)
+    return state, label_map, class_names
+
+def _preprocess_wav(y: np.ndarray) -> np.ndarray:
+    """학습 KWSDataset._preprocess와 동일: RMS norm + pre-emphasis (서버 추론 오탐 감소)."""
+    y = np.asarray(y, dtype=np.float32)
+    if len(y) == 0:
+        return y
+    rms = np.sqrt(np.mean(y ** 2)) + 1e-8
+    y = y * (RMS_NORM_TARGET / rms)
+    y = np.append(y[0], y[1:] - PRE_EMPHASIS * y[:-1]).astype(np.float32)
+    return y
+
+
+def _audio_to_mel(y: np.ndarray, sr: int) -> np.ndarray:
+    if sr != SR_MODEL:
+        y = librosa.resample(y.astype(np.float32), orig_sr=sr, target_sr=SR_MODEL)
+    y = _preprocess_wav(y)
+    S = librosa.feature.melspectrogram(
+        y=y, sr=SR_MODEL, n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP_LENGTH
+    )
+    S_db = np.clip(librosa.power_to_db(S, ref=1.0), -80.0, 0.0)
+    norm_mel = (S_db + 80.0) / 80.0
+    return norm_mel.astype(np.float32)
+
+def _mel_window_to_tensor(mel: np.ndarray, start_f: int, end_f: int) -> torch.Tensor:
+    chunk = mel[:, start_f:end_f]
+    n_f = chunk.shape[1]
+    if n_f >= MAX_TIME_FRAMES:
+        chunk = chunk[:, :MAX_TIME_FRAMES]
+    else:
+        chunk = np.pad(chunk, ((0, 0), (0, MAX_TIME_FRAMES - n_f)), mode="constant", constant_values=0)
+    t = torch.from_numpy(chunk).float().unsqueeze(0).unsqueeze(0).to(DEVICE)
+    return t
+
+def _frame_to_sec(frame_index: int) -> float:
+    return float(frame_index * HOP_LENGTH) / SR_MODEL
+
+class KWSLongInference:
+    def __init__(
+        self,
+        checkpoint_path: str,
+        window_sec: float = WINDOW_SEC,
+        step_sec: float = 0.25,
+        rms_threshold: float = 0.005, # 소리 크기 임계값 하향 조정
+    ) -> None:
+        self.checkpoint_path = checkpoint_path
+        self.window_sec = window_sec
+        self.step_sec = step_sec
+        self.rms_threshold = rms_threshold
+        
+        state, label_map, class_names = _load_checkpoint(checkpoint_path)
+        self.label_map = label_map
+        self.class_names = class_names
+        self.num_classes = len(class_names)
+        self.model = KWSModel(num_classes=self.num_classes).to(DEVICE)
+        self.model.load_state_dict(state, strict=True)
+        self.model.eval()
+        
+        self.window_frames = MAX_TIME_FRAMES
+        self.step_frames = max(1, int(round(step_sec * SR_MODEL / HOP_LENGTH)))
+
+    def _get_windows(self, total_frames: int) -> List[Tuple[int, int]]:
+        windows: List[Tuple[int, int]] = []
+        start = 0
+        while start + self.window_frames <= total_frames:
+            windows.append((start, start + self.window_frames))
+            start += self.step_frames
+        return windows
+
+    def predict_long(self, audio: Union[str, Tuple[int, np.ndarray]]) -> LongAudioResult:
+        if isinstance(audio, str):
+            y, sr = librosa.load(audio, sr=SR_MODEL)
+        else:
+            sr, y = audio
+            y = np.asarray(y, dtype=np.float32)
+            if sr != SR_MODEL:
+                y = librosa.resample(y, orig_sr=sr, target_sr=SR_MODEL)
+        
+        duration_sec = len(y) / SR_MODEL
+        mel = _audio_to_mel(y, SR_MODEL)
+        total_frames = mel.shape[1]
+        windows = self._get_windows(total_frames)
+
+        all_probs: List[np.ndarray] = []
+        window_results: List[WindowResult] = []
+
+        with torch.no_grad():
+            for start_f, end_f in windows:
+                # 윈도우 구간의 오디오 샘플 추출하여 RMS 계산
+                start_sample = int(start_f * HOP_LENGTH)
+                end_sample = int(end_f * HOP_LENGTH)
+                y_chunk = y[start_sample:end_sample]
+                rms = np.sqrt(np.mean(y_chunk**2)) if len(y_chunk) > 0 else 0
+                
+                start_sec = _frame_to_sec(start_f)
+                end_sec = _frame_to_sec(end_f)
+
+                if rms < self.rms_threshold:
+                    # [RMS 필터링] 소리가 너무 작으면 무조건 Normal로 처리
+                    probs = np.zeros(self.num_classes)
+                    probs[self.label_map.get("normal", 0)] = 1.0
+                    pred_id = self.label_map.get("normal", 0)
+                else:
+                    # 모델 추론
+                    t = _mel_window_to_tensor(mel, start_f, end_f)
+                    logits = self.model(t)
+                    probs = F.softmax(logits, dim=1).cpu().numpy()[0]
+                    pred_id = int(np.argmax(probs))
+
+                all_probs.append(probs)
+                conf = float(probs[pred_id])
+                probs_dict = {name: float(probs[i]) for i, name in enumerate(self.class_names)}
+                
+                window_results.append(
+                    WindowResult(
+                        start_sec=start_sec,
+                        end_sec=end_sec,
+                        pred_id=pred_id,
+                        class_name=self.class_names[pred_id],
+                        probs=probs_dict,
+                        confidence=conf,
+                    )
+                )
+
+        if not all_probs:
+            return LongAudioResult("normal", {}, 0, [], duration_sec)
+
+        # Max over windows 집계
+        stacked = np.array(all_probs)
+        max_per_class = np.max(stacked, axis=0)
+        agg_pred_id = int(np.argmax(max_per_class))
+        agg_label = self.class_names[agg_pred_id]
+        agg_probs = {name: float(max_per_class[i]) for i, name in enumerate(self.class_names)}
+
+        return LongAudioResult(agg_label, agg_probs, agg_pred_id, window_results, duration_sec)
+
+def run_long_inference(
+    wav_path: str,
+    checkpoint_path: str = "kws_final_model/best_kws_model.pth",
+    step_sec: float = 0.25,
+) -> LongAudioResult:
+    engine = KWSLongInference(checkpoint_path=checkpoint_path, step_sec=step_sec)
+    return engine.predict_long(wav_path)

kws_models_fpfix/best_kws_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77212523c6e434385d36c1bd2a972f8a49706af0125357b4f6070d221f5d3283
+size 1693516

kws_models_fpfix/training_config.json

@@ -0,0 +1,48 @@
+{
+  "manifest": "./dataset_hf/manifests/kws.jsonl",
+  "dataset_root": "./dataset_hf",
+  "data_dir": "./kws_data",
+  "output_dir": "./kws_models_fpfix",
+  "epochs": 80,
+  "lr": 0.0005,
+  "batch_size": 64,
+  "keywords": [
+    "help_me",
+    "save_me"
+  ],
+  "include_normal": true,
+  "normal_dir": "/home/dusen0528/LastResNet/data",
+  "trim_silence": false,
+  "rms_norm": 0.05,
+  "augment_noise_scale": 0.005,
+  "augment_shift": 0.1,
+  "augment_pitch": "0,2",
+  "augment_time_stretch": "0.9,1.1",
+  "spec_augment_on": true,
+  "fast_augment": true,
+  "mel_gpu": true,
+  "class_weight": "inverse",
+  "dropout": 0.3,
+  "weight_decay": 0.0001,
+  "early_stop_patience": 20,
+  "lr_scheduler_patience": 5,
+  "hf_repo": "",
+  "hf_token": "",
+  "wandb_project": "kws-fpfix",
+  "wandb_entity": "",
+  "amp": true,
+  "workers": 4,
+  "best_epoch": 51,
+  "val_recall_positive_min": 0.926829268292683,
+  "class_names": [
+    "normal",
+    "help_me",
+    "save_me"
+  ],
+  "label_map": {
+    "normal": 0,
+    "help_me": 1,
+    "save_me": 2
+  },
+  "normal_dir_added": 1000
+}

requirements.txt

@@ -0,0 +1,8 @@
+gradio>=6.5.1
+numpy
+librosa
+torch
+safetensors
+scikit-learn
+wandb
+tqdm

train_kws.py

@@ -0,0 +1,730 @@
+"""
+Keyword Spotting (KWS) 모델 학습 스크립트 - Production Optimized
+- WandB Sweep 연동 (augment_shift, noise, class_weight 등)
+- False Positive Rate (오탐률) 메트릭 추가
+- Audio Shift Augmentation 구현
+"""
+import os
+import json
+import random
+from datetime import datetime
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import librosa
+import argparse
+from pathlib import Path
+from sklearn.metrics import f1_score, precision_recall_fscore_support, accuracy_score, confusion_matrix
+from sklearn.model_selection import train_test_split
+import wandb
+from tqdm import tqdm
+
+# --- Global Config ---
+SR_MODEL = 16000
+N_MELS = 128
+N_FFT = 1024
+HOP_LENGTH = 512
+WINDOW_SEC = 1.5
+MAX_TIME_FRAMES = int(round((WINDOW_SEC * SR_MODEL - N_FFT) / HOP_LENGTH + 1))
+# 고정 입력 프레임에 필요한 샘플 수 (GPU 멜용 고정 길이)
+MAX_AUDIO_SAMPLES = (MAX_TIME_FRAMES - 1) * HOP_LENGTH + N_FFT
+
+def _compute_mel_fast(y: np.ndarray, sr: int) -> np.ndarray:
+    """멜 스펙트로그램 계산. torchaudio 사용 시 librosa보다 빠름, 동일 (n_mels, time) 0~1."""
+    try:
+        import torchaudio
+        t = torch.from_numpy(y).float().unsqueeze(0)
+        mel = torchaudio.functional.mel_spectrogram(
+            t, n_fft=N_FFT, hop_length=HOP_LENGTH, win_length=N_FFT,
+            f_min=0.0, f_max=float(sr) / 2, n_mels=N_MELS, power=2.0,
+        )
+        mel = mel.squeeze(0).numpy()
+        mel_db = np.clip(10.0 * np.log10(np.maximum(mel, 1e-10)), -80.0, 0.0)
+        return ((mel_db + 80.0) / 80.0).astype(np.float32)
+    except Exception:
+        S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP_LENGTH)
+        S_db = np.clip(librosa.power_to_db(S, ref=1.0), -80.0, 0.0)
+        return ((S_db + 80.0) / 80.0).astype(np.float32)
+
+def _load_audio_fast(path: str) -> tuple[np.ndarray, int]:
+    """torchaudio 우선 사용 (librosa보다 로딩 빠름), 실패 시 librosa 폴백."""
+    try:
+        import torchaudio
+        wav, sr = torchaudio.load(path)
+        if wav.shape[0] > 1:
+            wav = wav.mean(dim=0, keepdim=True)
+        if sr != SR_MODEL:
+            wav = torchaudio.functional.resample(wav, sr, SR_MODEL)
+        return wav.squeeze(0).numpy().astype(np.float32), SR_MODEL
+    except Exception:
+        y, sr = librosa.load(path, sr=SR_MODEL)
+        return y.astype(np.float32), sr
+
+def str2bool(v):
+    if isinstance(v, bool): return v
+    return v.lower() in ('yes', 'true', 't', 'y', '1')
+
+def _parse_range(s):
+    """'0.8,1.2' 같은 문자열을 (0.8, 1.2) 튜플로 변환"""
+    if isinstance(s, list): return tuple(s)
+    try: 
+        s = str(s).strip("[]").split(",")
+        return (float(s[0]), float(s[1]))
+    except: return (0, 0)
+
+class KWSDataset(Dataset):
+    def __init__(self, items, label_map=None, trim_silence=False, trim_top_db=30.0,
+                 rms_norm_target=0.05, pre_emphasis=0.97, augment_noise_scale=0.0,
+                 augment_time_stretch=(0.0, 0.0), augment_pitch=(0.0, 0.0),
+                 augment_shift=0.0, spec_augment=(0, 0), preprocess_cache=None, mel_cache=None,
+                 skip_heavy_augment=False, mel_gpu=False):
+        self.label_map = label_map
+        self.items = items
+        self.trim_silence = trim_silence
+        self.trim_top_db = trim_top_db
+        self.rms_norm_target = rms_norm_target
+        self.pre_emphasis = pre_emphasis
+        self.preprocess_cache = preprocess_cache
+        self.mel_cache = mel_cache if not mel_gpu else None  # mel_gpu면 Dataset에서 멜 안 함
+        self.skip_heavy_augment = skip_heavy_augment
+        self.mel_gpu = mel_gpu
+
+        # Augmentations
+        self.augment_noise_scale = augment_noise_scale
+        self.augment_time_stretch = augment_time_stretch
+        self.augment_pitch = augment_pitch
+        self.augment_shift = augment_shift
+        self.spec_augment = spec_augment
+
+    def __len__(self): return len(self.items)
+
+    def _preprocess(self, y):
+        # 1. Silence Trim (옵션: 오탐 방지를 위해 보통 False 권장)
+        if self.trim_silence and len(y) > 0:
+            y, _ = librosa.effects.trim(y, top_db=self.trim_top_db)
+        if len(y) == 0: y = np.zeros(1024, dtype=np.float32)
+        
+        # 2. RMS Norm (볼륨 정규화)
+        if self.rms_norm_target > 0:
+            rms = np.sqrt(np.mean(y ** 2)) + 1e-8
+            y = y * (self.rms_norm_target / rms)
+            
+        # 3. Pre-emphasis (고주파 강조)
+        if self.pre_emphasis != 0:
+            y = np.append(y[0], y[1:] - self.pre_emphasis * y[:-1]).astype(np.float32)
+        return y
+
+    def _apply_augment(self, y, sr):
+        # A. Time Stretch (무거움, --fast-augment 시 생략)
+        if not self.skip_heavy_augment and self.augment_time_stretch[0] != self.augment_time_stretch[1]:
+            rate = random.uniform(self.augment_time_stretch[0], self.augment_time_stretch[1])
+            if abs(rate - 1.0) > 0.01:
+                y = librosa.effects.time_stretch(y, rate=rate)
+
+        # B. Pitch Shift (무거움, --fast-augment 시 생략)
+        if not self.skip_heavy_augment and self.augment_pitch[0] != self.augment_pitch[1]:
+            n_steps = random.uniform(self.augment_pitch[0], self.augment_pitch[1])
+            if abs(n_steps) > 0.01:
+                y = librosa.effects.pitch_shift(y, sr=sr, n_steps=n_steps)
+
+        # C. Time Shift (위치 이동) - [New] YAML 대응
+        # 오디오를 좌우로 밀고 빈 공간은 0(침묵)으로 채움 (np.roll은 순환이라 비추천)
+        if self.augment_shift > 0:
+            shift_sec = random.uniform(-self.augment_shift, self.augment_shift)
+            shift_samples = int(shift_sec * sr)
+            if abs(shift_samples) > 0:
+                y_shifted = np.zeros_like(y)
+                if shift_samples > 0: # 오른쪽으로 밀기 (앞에 침묵)
+                    if shift_samples < len(y):
+                        y_shifted[shift_samples:] = y[:-shift_samples]
+                else: # 왼쪽으로 밀기 (뒤에 침묵)
+                    shift_samples = abs(shift_samples)
+                    if shift_samples < len(y):
+                        y_shifted[:-shift_samples] = y[shift_samples:]
+                y = y_shifted
+
+        # D. Noise Injection (소음) - [Updated] 정규화 이후에 넣어야 일관됨
+        if self.augment_noise_scale > 0:
+            rms = np.sqrt(np.mean(y ** 2)) + 1e-8
+            # White Noise 생성 (실제 환경음 Mix가 더 좋지만, 현재는 Gaussian으로 대체)
+            noise = np.random.randn(len(y)).astype(np.float32)
+            y = y + self.augment_noise_scale * rms * noise
+            
+        return y
+
+    def _apply_spec_augment(self, mel):
+        t_max, f_max = self.spec_augment
+        if t_max <= 0 and f_max <= 0: return mel
+        n_mels, T = mel.shape
+        # Time Masking
+        if t_max > 0 and T > t_max:
+            t0 = random.randint(0, T - t_max)
+            mel[:, t0 : t0 + t_max] = mel.mean() # 평균값으로 마스킹
+        # Freq Masking
+        if f_max > 0 and n_mels > f_max:
+            f0 = random.randint(0, n_mels - f_max)
+            mel[f0 : f0 + f_max, :] = mel.mean()
+        return mel
+
+    def __getitem__(self, idx):
+        item = self.items[idx]
+        path = os.path.normpath(item['wav_path'])
+        try:
+            # mel_gpu가 아니고 검증용 멜 캐시 히트면 그대로 반환
+            if not self.mel_gpu and self.mel_cache is not None and path in self.mel_cache:
+                return self.mel_cache[path].clone(), item['label_id']
+
+            # 1. Load + 전처리
+            if self.preprocess_cache is not None:
+                if path in self.preprocess_cache:
+                    y, sr = self.preprocess_cache[path]
+                else:
+                    y, sr = _load_audio_fast(item['wav_path'])
+                    y = self._preprocess(y)
+                    self.preprocess_cache[path] = (y.copy(), sr)
+            else:
+                y, sr = _load_audio_fast(item['wav_path'])
+                y = self._preprocess(y)
+
+            # 2. Augment
+            y = self._apply_augment(y, sr)
+
+            if self.mel_gpu:
+                # 3a. GPU 멜 모드: 파형만 패딩/자르기 후 반환 (멜은 학습 루프에서 GPU 배치 계산)
+                if len(y) > MAX_AUDIO_SAMPLES:
+                    y = y[:MAX_AUDIO_SAMPLES]
+                else:
+                    y = np.pad(y, (0, MAX_AUDIO_SAMPLES - len(y)), mode="constant", constant_values=0)
+                return torch.from_numpy(y.astype(np.float32)), item['label_id']
+
+            # 3b. CPU 멜 모드: 기존처럼 멜 계산 후 반환
+            norm_mel = _compute_mel_fast(y, sr)
+            norm_mel = self._apply_spec_augment(norm_mel.copy())
+            if norm_mel.shape[1] > MAX_TIME_FRAMES:
+                norm_mel = norm_mel[:, :MAX_TIME_FRAMES]
+            else:
+                norm_mel = np.pad(norm_mel, ((0, 0), (0, MAX_TIME_FRAMES - norm_mel.shape[1])))
+            out = torch.from_numpy(norm_mel).float().unsqueeze(0)
+            if self.mel_cache is not None:
+                self.mel_cache[path] = out.clone()
+            return out, item['label_id']
+        except Exception as e:
+            print(f"⚠️ Load Error ({item['wav_path']}): {e}")
+            if self.mel_gpu:
+                return torch.zeros(MAX_AUDIO_SAMPLES, dtype=torch.float32), item['label_id']
+            return torch.zeros(1, N_MELS, MAX_TIME_FRAMES), item['label_id']
+
+class KWSModel(nn.Module):
+    def __init__(self, num_classes=3, n_mels=128, dropout=0.0):
+        super(KWSModel, self).__init__()
+        self.conv1 = nn.Conv1d(n_mels, 64, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm1d(64)
+        self.relu = nn.ReLU()
+        # Stride를 활용해 시간 차원 축소
+        self.layer1 = self._make_layer(64, 64, 1)
+        self.layer2 = self._make_layer(64, 128, 2)
+        self.layer3 = self._make_layer(128, 256, 2)
+        self.adaptive_pool = nn.AdaptiveAvgPool1d(1)
+        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
+        self.fc = nn.Linear(256, num_classes)
+
+    def _make_layer(self, in_c, out_c, stride):
+        return nn.Sequential(
+            nn.Conv1d(in_c, out_c, 3, stride, 1, bias=False), nn.BatchNorm1d(out_c), nn.ReLU(),
+            nn.Conv1d(out_c, out_c, 3, 1, 1, bias=False), nn.BatchNorm1d(out_c), nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = x.squeeze(1) # [B, 1, F, T] -> [B, F, T]
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.layer3(self.layer2(self.layer1(x)))
+        x = self.adaptive_pool(x).flatten(1)
+        return self.fc(self.dropout(x))
+
+
+class MelExtractorGPU(nn.Module):
+    """GPU에서 배치 단위로 멜 스펙트로그램 + (옵션) Spec Augment. Dataset은 파형만 반환하면 됨."""
+    def __init__(self, spec_augment: tuple[int, int] = (0, 0)):
+        super().__init__()
+        self.spec_augment = spec_augment
+        try:
+            import torchaudio
+            self._mel_fn = torchaudio.transforms.MelSpectrogram(
+                sample_rate=SR_MODEL,
+                n_fft=N_FFT,
+                win_length=N_FFT,
+                hop_length=HOP_LENGTH,
+                f_min=0.0,
+                f_max=float(SR_MODEL) / 2,
+                n_mels=N_MELS,
+                power=2.0,
+            )
+        except Exception:
+            self._mel_fn = None
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        # waveform: (B, samples)
+        if self._mel_fn is None:
+            raise RuntimeError("MelExtractorGPU requires torchaudio")
+        mel = self._mel_fn(waveform)
+        mel_db = 10.0 * torch.log10(torch.clamp(mel, min=1e-10))
+        mel_db = torch.clamp(mel_db, -80.0, 0.0)
+        norm_mel = (mel_db + 80.0) / 80.0
+        # (B, n_mels, time) -> pad/crop to MAX_TIME_FRAMES
+        T = norm_mel.shape[2]
+        if T > MAX_TIME_FRAMES:
+            norm_mel = norm_mel[:, :, :MAX_TIME_FRAMES]
+        elif T < MAX_TIME_FRAMES:
+            norm_mel = torch.nn.functional.pad(norm_mel, (0, MAX_TIME_FRAMES - T))
+        if self.training and self.spec_augment[0] > 0 and self.spec_augment[1] > 0:
+            norm_mel = self._spec_augment(norm_mel)
+        return norm_mel.unsqueeze(1)  # (B, 1, N_MELS, MAX_TIME_FRAMES)
+
+    def _spec_augment(self, mel: torch.Tensor) -> torch.Tensor:
+        B, F, T = mel.shape
+        t_max, f_max = self.spec_augment
+        if t_max > 0 and T > t_max:
+            t0 = torch.randint(0, T - t_max + 1, (B,), device=mel.device)
+            for i in range(B):
+                mel[i, :, t0[i] : t0[i] + t_max] = mel[i].mean()
+        if f_max > 0 and F > f_max:
+            f0 = torch.randint(0, F - f_max + 1, (B,), device=mel.device)
+            for i in range(B):
+                mel[i, f0[i] : f0[i] + f_max, :] = mel[i].mean()
+        return mel
+
+
+def _serializable_config(args, best_epoch, best_metric, class_names, label_map, normal_dir_added=0):
+    out = {}
+    for k, v in vars(args).items():
+        if v is None or (isinstance(v, float) and (v != v)): continue
+        if isinstance(v, (str, int, float, bool)) or (isinstance(v, (list, tuple)) and all(isinstance(x, (str, int, float, bool)) for x in (v or []))):
+            out[k] = v
+        else: out[k] = str(v)
+    out.update({
+        "best_epoch": best_epoch,
+        "val_recall_positive_min": best_metric,
+        "class_names": class_names,
+        "label_map": label_map,
+        "normal_dir_added": normal_dir_added
+    })
+    return out
+
+def main():
+    parser = argparse.ArgumentParser()
+    # Path
+    parser.add_argument('--manifest', type=str, default='')
+    parser.add_argument('--dataset-root', type=str, default='')
+    parser.add_argument('--data-dir', type=str, default='./kws_data')
+    parser.add_argument('--output-dir', type=str, default='./kws_models')
+    
+    # Train
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--batch-size', type=int, default=32)
+    parser.add_argument('--keywords', nargs='+', default=['help_me', 'save_me'])
+    parser.add_argument('--include-normal', action='store_true')
+    parser.add_argument('--normal-dir', type=str, default='', help='일반(Normal) 데이터 추가 경로')
+    
+    # Augment
+    parser.add_argument('--trim_silence', type=str2bool, default=False)
+    parser.add_argument('--rms_norm', type=float, default=0.05)
+    parser.add_argument('--augment_noise_scale', type=float, default=0.0)
+    parser.add_argument('--augment_shift', type=float, default=0.0, help='[New] Time Shift Augmentation (sec)')
+    parser.add_argument('--augment_pitch', type=str, default='0,0')
+    parser.add_argument('--augment_time_stretch', type=str, default='0,0')
+    parser.add_argument('--spec_augment_on', type=str2bool, default=False)
+    parser.add_argument('--fast-augment', action='store_true', help='time_stretch/pitch_shift 생략 → 워커 CPU 부하 감소, GPU 활용도 상승 (sweep 병목 시 사용)')
+    parser.add_argument('--mel-gpu', action='store_true', help='멜 스펙트로그램을 GPU에서 배치 계산 → 워커 CPU 부하 감소, 학습 가속')
+    
+    # Regularization (과적합 완화)
+    parser.add_argument('--class_weight', type=str, default='inverse')
+    parser.add_argument('--dropout', type=float, default=0.0)
+    parser.add_argument('--weight_decay', type=float, default=1e-4, help='L2 정규화 (Adam). 0이면 비활성화')
+    parser.add_argument('--early_stop_patience', type=int, default=10)
+    parser.add_argument('--lr_scheduler_patience', type=int, default=0, help='val_loss 개선 없을 때 LR 감소 대기 에폭. 0이면 비활성화')
+    
+    # HF / WandB (HF 기본: dusen0528/kws)
+    parser.add_argument('--hf-repo', type=str, default='', help='Hugging Face 모델 repo (기본: dusen0528/kws)')
+    parser.add_argument('--hf-token', type=str, default='')
+    parser.add_argument('--wandb-project', type=str, default='kws')
+    parser.add_argument('--wandb-entity', type=str, default='')
+    parser.add_argument('--amp', type=str2bool, default=True, help='GPU에서 혼합정밀(AMP) 사용 → 학습 속도 향상')
+    parser.add_argument('--workers', type=int, default=4, help='DataLoader worker 수 (0이면 메인만)')
+    
+    args, _ = parser.parse_known_args()
+
+    # --- 1. Label Map & Data Loading ---
+    label_map = {}
+    # Normal이 있으면 0번으로 고정 (관례)
+    if args.include_normal or (args.normal_dir and os.path.isdir(args.normal_dir)):
+        label_map['normal'] = 0
+    for kw in args.keywords: label_map[kw] = len(label_map)
+    num_classes = len(label_map)
+    class_names = [k for k, v in sorted(label_map.items(), key=lambda x: x[1])]
+
+    records = []
+    # A. Manifest Load
+    if args.manifest and os.path.exists(args.manifest):
+        root = Path(args.dataset_root)
+        with open(args.manifest, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    obj = json.loads(line)
+                    lbl = obj.get("label")
+                    if lbl in label_map:
+                        ap = obj["audio_path"]
+                        full_p = ap if os.path.isabs(ap) else str(root / ap)
+                        if os.path.exists(full_p):
+                            records.append({"wav_path": full_p, "label": lbl, "label_id": label_map[lbl]})
+                except: pass
+    
+    # B. Folder Load (Fallback)
+    if not records:
+        for ln, lid in label_map.items():
+            dp = os.path.join(args.data_dir, ln)
+            if os.path.isdir(dp):
+                for f in os.listdir(dp):
+                    if f.lower().endswith(('.wav', '.mp3')):
+                        records.append({"wav_path": os.path.join(dp, f), "label": ln, "label_id": lid})
+
+    # C. Normal Dir Augmentation
+    normal_dir_added = 0
+    if args.normal_dir and os.path.isdir(args.normal_dir) and 'normal' in label_map:
+        n_before = len(records)
+        for f in os.listdir(args.normal_dir):
+            if f.lower().endswith(('.wav', '.mp3')):
+                full_p = os.path.join(args.normal_dir, f)
+                if os.path.isfile(full_p):
+                    records.append({"wav_path": full_p, "label": "normal", "label_id": label_map["normal"]})
+        normal_dir_added = len(records) - n_before
+
+    if not records: print("❌ 데이터 없음"); return
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    hf_repo = (args.hf_repo or os.environ.get("HF_REPO", "")).strip() or "dusen0528/kws"
+    hf_token = args.hf_token or os.environ.get("HF_TOKEN", "")
+
+    # --- 2. WandB Init ---
+    run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
+    if os.environ.get("WANDB_SWEEP_ID"):
+        wandb.init(config=vars(args), name=run_name)
+        # Sweep 파라미터가 args를 덮어씀
+        for k, v in wandb.config.items():
+            setattr(args, k, v)
+    else:
+        wandb.init(config=vars(args), project=args.wandb_project, entity=args.wandb_entity or None, name=run_name)
+
+    # --- 3. Dataset & Loader ---
+    stratify_ids = [r['label_id'] for r in records]
+    try:
+        train_items, val_items = train_test_split(records, test_size=0.2, random_state=42, stratify=stratify_ids)
+    except ValueError:
+        train_items, val_items = train_test_split(records, test_size=0.2, random_state=42)
+    
+    preprocess_cache: dict = {}
+    fast_augment = getattr(args, 'fast_augment', False)
+    mel_gpu = getattr(args, "mel_gpu", False)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if mel_gpu and device.type != "cuda":
+        mel_gpu = False
+        print("   ⚠️ --mel-gpu는 CUDA에서만 동작, 비활성화함")
+    train_ds = KWSDataset(train_items, label_map,
+                          trim_silence=args.trim_silence,
+                          rms_norm_target=args.rms_norm,
+                          augment_noise_scale=args.augment_noise_scale,
+                          augment_pitch=_parse_range(args.augment_pitch),
+                          augment_time_stretch=_parse_range(args.augment_time_stretch),
+                          augment_shift=args.augment_shift,
+                          spec_augment=(15, 10) if args.spec_augment_on else (0, 0),
+                          preprocess_cache=preprocess_cache,
+                          skip_heavy_augment=fast_augment,
+                          mel_gpu=mel_gpu)
+    
+    val_mel_cache: dict = {} if not mel_gpu else {}
+    val_ds = KWSDataset(val_items, label_map,
+                        trim_silence=args.trim_silence,
+                        rms_norm_target=args.rms_norm,
+                        preprocess_cache=preprocess_cache,
+                        mel_cache=val_mel_cache,
+                        mel_gpu=mel_gpu)
+    
+    use_amp = args.amp and device.type == 'cuda'
+    num_workers = args.workers if device.type == 'cuda' else 0
+    pin = (device.type == 'cuda')
+    loader_kw = dict(batch_size=args.batch_size, pin_memory=pin, num_workers=num_workers)
+    if num_workers > 0:
+        loader_kw["persistent_workers"] = True
+        loader_kw["prefetch_factor"] = 8
+    train_loader = DataLoader(train_ds, shuffle=True, **loader_kw)
+    val_loader = DataLoader(val_ds, **loader_kw)
+    if use_amp:
+        print("⚡ AMP(혼합정밀) 사용")
+    
+    # --- 4. Model & Loss ---
+    model = KWSModel(num_classes=num_classes, dropout=args.dropout).to(device)
+    mel_extractor = None
+    if mel_gpu and device.type == "cuda":
+        mel_extractor = MelExtractorGPU(
+            spec_augment=(15, 10) if args.spec_augment_on else (0, 0)
+        ).to(device)
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=args.lr,
+        weight_decay=getattr(args, 'weight_decay', 0.0),
+    )
+    scheduler = None
+    if getattr(args, 'lr_scheduler_patience', 0) > 0:
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode='min', factor=0.5, patience=args.lr_scheduler_patience, min_lr=1e-6
+        )
+        print(f"   📉 LR Scheduler: val_loss {args.lr_scheduler_patience}ep 개선 없으면 LR×0.5")
+    
+    class_weights = None
+    if args.class_weight == 'inverse':
+        from collections import Counter
+        counts = Counter([r['label_id'] for r in train_items])
+        # N_samples / (N_classes * count)
+        weights = [len(train_items) / (num_classes * (counts.get(i, 0) + 1)) for i in range(num_classes)]
+        class_weights = torch.tensor(weights, dtype=torch.float32).to(device)
+        print(f"⚖️ Class Weights: {class_weights.tolist()}")
+
+    criterion = nn.CrossEntropyLoss(weight=class_weights)
+    scaler = torch.amp.GradScaler("cuda") if use_amp else None
+
+    # --- 5. Training Loop ---
+    best_rec_min = 0.0
+    best_fpr = 1.0
+    best_epoch = 0
+    patience = 0
+    
+    print(f"🚀 학습 시작 (데이터: {len(records)}개, 클래스: {class_names})")
+    print(f"   workers={num_workers} batch={args.batch_size} | GPU 낮고 pt_data_worker 100%면: --batch-size 64 --workers 8, --fast-augment, --mel-gpu")
+    if fast_augment:
+        print(f"   ⚡ --fast-augment: time_stretch/pitch_shift 생략 → 워커 부하 감소")
+    if mel_gpu and mel_extractor is not None:
+        print(f"   ⚡ --mel-gpu: 멜 스펙트로그램을 GPU에서 배치 계산 → 워커 CPU 부하 감소")
+
+    pbar_epoch = tqdm(range(1, args.epochs + 1), desc="Epoch", unit="ep")
+    for e in pbar_epoch:
+        # A. Train
+        model.train()
+        t_loss = 0; t_corr = 0; t_total = 0
+        train_pbar = tqdm(train_loader, desc=f"Ep {e}", leave=False, unit="batch")
+        for m, l in train_pbar:
+            m, l = m.to(device, non_blocking=True), l.to(device, non_blocking=True)
+            if mel_extractor is not None:
+                mel_extractor.train()
+                m = mel_extractor(m)
+            optimizer.zero_grad()
+            if use_amp and scaler is not None:
+                with torch.amp.autocast("cuda"):
+                    out = model(m)
+                    loss = criterion(out, l)
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                out = model(m)
+                loss = criterion(out, l)
+                loss.backward()
+                optimizer.step()
+            t_loss += loss.item()
+            _, p = torch.max(out, 1)
+            t_corr += (p == l).sum().item()
+            t_total += l.size(0)
+            train_pbar.set_postfix(loss=f"{t_loss / max(1, train_pbar.n):.3f}")
+        
+        train_loss = t_loss / len(train_loader)
+        train_acc = 100 * t_corr / t_total if t_total else 0
+
+        # B. Validation
+        model.eval()
+        v_loss = 0; v_preds, v_labels = [], []
+        with torch.no_grad():
+            for m, l in tqdm(val_loader, desc="Val", leave=False, unit="batch"):
+                m, l = m.to(device, non_blocking=True), l.to(device, non_blocking=True)
+                if mel_extractor is not None:
+                    mel_extractor.eval()
+                    with torch.no_grad():
+                        m = mel_extractor(m)
+                if use_amp:
+                    with torch.amp.autocast("cuda"):
+                        out = model(m)
+                        loss = criterion(out, l)
+                else:
+                    out = model(m)
+                    loss = criterion(out, l)
+                v_loss += loss.item()
+                _, p = torch.max(out, 1)
+                v_preds.extend(p.cpu().numpy())
+                v_labels.extend(l.cpu().numpy())
+        
+        val_loss = v_loss / len(val_loader) if len(val_loader) else 0
+        val_acc = 100 * accuracy_score(v_labels, v_preds) if v_labels else 0
+
+        # C. Metrics (Recall per class & False Positive Rate)
+        labels_idx = list(range(num_classes))
+        prec, rec, f1, _ = precision_recall_fscore_support(
+            v_labels, v_preds, labels=labels_idx, average=None, zero_division=0
+        )
+        rec = np.atleast_1d(rec).astype(float)
+        prec = np.atleast_1d(prec).astype(float)
+        f1 = np.atleast_1d(f1).astype(float)
+        if len(rec) < num_classes:
+            rec = np.pad(rec, (0, num_classes - len(rec)), constant_values=0.0)
+            prec = np.pad(prec, (0, num_classes - len(prec)), constant_values=0.0)
+            f1 = np.pad(f1, (0, num_classes - len(f1)), constant_values=0.0)
+        
+        log_dict = {
+            "epoch": e,
+            "train/loss": train_loss, "train/acc": train_acc,
+            "val/loss": val_loss, "val/acc": val_acc,
+            "val/f1_macro": f1_score(v_labels, v_preds, average='macro', zero_division=0),
+            "val/train_loss_gap": val_loss - train_loss,  # >0 이면 val이 train보다 나쁨 → 오버피팅 의심
+            "val/train_acc_gap": train_acc - val_acc,      # >0 이면 train이 val보다 좋음 → 오버피팅 의심
+        }
+        if scheduler is not None:
+            log_dict["train/lr"] = optimizer.param_groups[0]["lr"]
+        
+        pos_rec = []
+        for i, name in enumerate(class_names):
+            safe_name = name.replace("/", "_")
+            log_dict[f"val/recall_{safe_name}"] = rec[i]
+            log_dict[f"val/f1_{safe_name}"] = f1[i]
+            if name in ['help_me', 'save_me']:
+                pos_rec.append(rec[i])
+        
+        # [Metric 1] Recall Positive Min (목표: 미탐 방지)
+        rec_min = min(pos_rec) if pos_rec else 0.0
+        log_dict["val/recall_positive_min"] = rec_min
+
+        # [Metric 2] False Positive Rate (목표: 오탐 방지)
+        # Normal 데이터(실제 0번)가 들어왔는데 -> 0번이 아니라고(Positive라고) 예측한 비율
+        if 'normal' in label_map:
+            norm_idx = label_map['normal']
+            norm_mask = [i for i, x in enumerate(v_labels) if x == norm_idx]
+            if norm_mask:
+                norm_preds = [v_preds[i] for i in norm_mask]
+                false_alarms = sum(1 for p in norm_preds if p != norm_idx)
+                fpr = false_alarms / len(norm_mask)
+            else:
+                fpr = 0.0
+            log_dict["val/false_positive_rate"] = fpr
+        
+        wandb.log(log_dict)
+        
+        current_fpr = log_dict.get("val/false_positive_rate", 1.0)
+        loss_gap = val_loss - train_loss
+        acc_gap = train_acc - val_acc
+        if scheduler is not None:
+            scheduler.step(val_loss)
+        pbar_epoch.set_postfix(loss=f"{train_loss:.3f}", rec_min=f"{rec_min:.2f}", fpr=f"{current_fpr:.2f}")
+        print(f"[Ep {e}] Loss: {train_loss:.4f} | RecMin: {rec_min:.2f} | FPR: {current_fpr:.2f}")
+        # 과적합 경고: 기준을 강하게 해서 매 에폭 뜨지 않게 함 (loss_gap 0.5 이상 또는 acc_gap 15% 이상일 때만)
+        if e > 5 and (loss_gap > 0.5 or acc_gap > 15.0):
+            print(f"    ⚠️ 오버피팅 가능성 (val_loss - train_loss = {loss_gap:.3f}, train_acc - val_acc = {acc_gap:.1f}%)")
+        if e == 10 and train_loss > 0.5 and val_loss > 0.5:
+            print(f"    💡 언더피팅 가능성 (Ep 10인데 loss 둘 다 높음 → epoch/모델/학습률 검토)")
+
+        # D. Checkpoint
+        # 저장 조건 1: Recall Min이 더 높으면 무조건 저장 (미탐 방지 최우선)
+        # 저장 조건 2: Recall Min이 같으면, 오탐률(FPR)이 더 낮은 모델 저장
+        is_best = False
+        if rec_min > best_rec_min:
+            is_best = True
+        elif rec_min == best_rec_min and rec_min > 0:
+            if current_fpr < best_fpr:
+                is_best = True
+                print(f"✨ Recall 동점({rec_min:.2f})이나 FPR 개선됨 ({best_fpr:.2f} -> {current_fpr:.2f})")
+
+        if is_best:
+            best_rec_min = rec_min
+            best_fpr = current_fpr
+            best_epoch = e
+            patience = 0
+            from checkpoint_io import save_checkpoint
+            save_checkpoint(
+                args.output_dir,
+                model.state_dict(),
+                _serializable_config(args, best_epoch, best_rec_min, class_names, label_map, normal_dir_added),
+                also_save_pth=False,
+            )
+            print(f"💾 Best 저장 | Ep {e} | RecMin {best_rec_min:.2f} | FPR {best_fpr:.2f} | → best_kws_model.safetensors")
+        else:
+            patience += 1
+            if args.early_stop_patience > 0 and patience >= args.early_stop_patience:
+                print("🛑 Early Stopping")
+                break
+
+    # 학습 종료
+    if getattr(wandb, "run", None) is not None:
+        wandb.run.summary["val/recall_positive_min"] = best_rec_min
+        wandb.run.summary["best_epoch"] = best_epoch
+        wandb.run.summary["best_fpr"] = best_fpr
+
+    # --- HF: sweep/학습 종료 후 업로드 (기준 날짜 폴더 + 파라미터·성능) ---
+    if hf_repo:
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi(token=hf_token or None)
+            api.create_repo(hf_repo, repo_type="model", exist_ok=True)
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+            folder = f"run_{ts}"
+            for f in ["best_kws_model.safetensors", "best_kws_model.pth", "training_config.json"]:
+                p = os.path.join(args.output_dir, f)
+                if os.path.exists(p):
+                    api.upload_file(path_or_fileobj=p, path_in_repo=f"{folder}/{f}", repo_id=hf_repo, repo_type="model")
+            metrics = {
+                "uploaded_at": datetime.now().isoformat(),
+                "type": "run_post_training",
+                "best_epoch": best_epoch,
+                "val_recall_positive_min": best_rec_min,
+                "best_fpr": best_fpr,
+            }
+            if getattr(wandb, "run", None) is not None:
+                run = wandb.run
+                metrics["wandb_run_id"] = getattr(run, "id", "") or ""
+                metrics["wandb_run_name"] = getattr(run, "name", "") or ""
+                metrics["wandb_run_url"] = getattr(run, "url", "") or ""
+            metrics_path = os.path.join(args.output_dir, "metrics.json")
+            with open(metrics_path, "w", encoding="utf-8") as mf:
+                json.dump(metrics, mf, indent=2, ensure_ascii=False)
+            api.upload_file(path_or_fileobj=metrics_path, path_in_repo=f"{folder}/metrics.json", repo_id=hf_repo, repo_type="model")
+            # best/ = 모든 run 중 진짜 최고만 유지. 기존 best 지표와 비교해 더 좋을 때만 덮어쓰기
+            should_update_best = True
+            existing_rec, existing_fpr = 0.0, 1.0
+            try:
+                from huggingface_hub import hf_hub_download
+                path = hf_hub_download(
+                    repo_id=hf_repo,
+                    filename="best/metrics.json",
+                    repo_type="model",
+                    token=hf_token or None,
+                )
+                with open(path, "r", encoding="utf-8") as f:
+                    existing = json.load(f)
+                existing_rec = float(existing.get("val_recall_positive_min", 0.0))
+                existing_fpr = float(existing.get("best_fpr", 1.0))
+                # 더 좋음 = rec_min 더 높음, 또는 rec_min 동점이면 FPR 더 낮음
+                if best_rec_min > existing_rec or (best_rec_min == existing_rec and best_fpr < existing_fpr):
+                    should_update_best = True
+                else:
+                    should_update_best = False
+            except Exception:
+                # best/metrics.json 없음 (첫 run) → best 갱신
+                should_update_best = True
+            if should_update_best:
+                for f in ["best_kws_model.safetensors", "best_kws_model.pth", "training_config.json"]:
+                    p = os.path.join(args.output_dir, f)
+                    if os.path.exists(p):
+                        api.upload_file(path_or_fileobj=p, path_in_repo=f"best/{f}", repo_id=hf_repo, repo_type="model")
+                api.upload_file(path_or_fileobj=metrics_path, path_in_repo="best/metrics.json", repo_id=hf_repo, repo_type="model")
+                print(f"[KWS] HF 업로드 완료: {hf_repo} -> {folder}/ (히스토리) + best/ (전체 run 중 최고로 갱신)")
+            else:
+                print(f"[KWS] HF 업로드 완료: {hf_repo} -> {folder}/ (히스토리). best/ 미갱신: 이번 rec_min={best_rec_min:.2f} fpr={best_fpr:.2f} vs 기존 rec_min={existing_rec:.2f} fpr={existing_fpr:.2f}")
+        except Exception as e:
+            print(f"[KWS] HF Upload Error: {e}")
+
+if __name__ == "__main__": main()