feat: add dataset statistics and class balance visualizations to enhance analysis capabilities

app/training/visualize_results.py

@@ -556,6 +556,387 @@ def _load_features_with_names(
     return X, y, feature_cols
 
 
+def _load_manifest(manifest_csv: Path) -> list[dict]:
+    """Load manifest.csv rows as dicts."""
+    if not manifest_csv.exists():
+        return []
+    with open(manifest_csv, "r", encoding="utf-8") as f:
+        return list(csv.DictReader(f))
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Figure 9: Dataset Statistics
+# ═══════════════════════════════════════════════════════════════════════
+
+def plot_dataset_statistics(
+    manifest_csv: Path,
+    output_dir: Path,
+) -> None:
+    """Dataset composition: class distribution, source distribution, duration."""
+    rows = _load_manifest(manifest_csv)
+    if not rows:
+        print("  Skipping dataset stats — no manifest")
+        return
+
+    from collections import Counter
+
+    labels = [r.get("label", "unknown") for r in rows]
+    sources = [r.get("generator", "unknown") for r in rows]
+    durations = []
+    for r in rows:
+        try:
+            durations.append(float(r.get("duration_sec", 0)))
+        except ValueError:
+            continue
+
+    label_counts = Counter(labels)
+    source_counts = Counter(sources)
+
+    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
+
+    # (a) Class distribution
+    ax = axes[0]
+    colors = [AURIS_BLUE if l == "human" else AURIS_RED for l in label_counts.keys()]
+    bars = ax.bar(label_counts.keys(), label_counts.values(),
+                  color=colors, edgecolor="black", linewidth=0.5)
+    for bar, val in zip(bars, label_counts.values()):
+        pct = val / sum(label_counts.values()) * 100
+        ax.text(bar.get_x() + bar.get_width() / 2,
+                bar.get_height() + max(label_counts.values()) * 0.02,
+                f"{val:,}\n({pct:.1f}%)",
+                ha="center", va="bottom", fontsize=10)
+    ax.set_title("Class Distribution")
+    ax.set_ylabel("Samples")
+    ax.set_ylim([0, max(label_counts.values()) * 1.18])
+
+    # (b) Source distribution
+    ax = axes[1]
+    sorted_sources = sorted(source_counts.items(), key=lambda x: x[1], reverse=True)
+    src_names = [s[0] for s in sorted_sources]
+    src_values = [s[1] for s in sorted_sources]
+    bars = ax.barh(src_names, src_values, color=plt.cm.tab10.colors[:len(src_names)])
+    for bar, val in zip(bars, src_values):
+        ax.text(bar.get_width() + max(src_values) * 0.01,
+                bar.get_y() + bar.get_height() / 2,
+                f"{val:,}", ha="left", va="center", fontsize=9)
+    ax.set_title("Dataset Source Distribution")
+    ax.set_xlabel("Samples")
+    ax.set_xlim([0, max(src_values) * 1.15])
+    ax.invert_yaxis()
+
+    # (c) Duration histogram
+    ax = axes[2]
+    if durations:
+        ax.hist(durations, bins=40, color=AURIS_BLUE, edgecolor="white", alpha=0.8)
+        ax.axvline(np.median(durations), color="red", linestyle="--",
+                   label=f"Median: {np.median(durations):.1f}s")
+        ax.axvline(np.mean(durations), color="orange", linestyle="--",
+                   label=f"Mean: {np.mean(durations):.1f}s")
+        ax.legend()
+    ax.set_title("Audio Duration Distribution")
+    ax.set_xlabel("Duration (seconds)")
+    ax.set_ylabel("Count")
+
+    fig.suptitle(
+        f"AURIS Dataset — {sum(label_counts.values()):,} samples",
+        fontsize=14, y=1.02,
+    )
+    fig.tight_layout()
+    _save_fig(fig, output_dir, "fig9_dataset_statistics")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Figure 10: Per-Source Class Balance
+# ═══════════════════════════════════════════════════════════════════════
+
+def plot_per_source_balance(
+    manifest_csv: Path,
+    output_dir: Path,
+) -> None:
+    """Stacked bar: AI vs Human count per source."""
+    rows = _load_manifest(manifest_csv)
+    if not rows:
+        return
+
+    from collections import defaultdict
+    source_label_counts: dict[str, dict[str, int]] = defaultdict(
+        lambda: {"ai": 0, "human": 0}
+    )
+    for r in rows:
+        src = r.get("generator", "unknown")
+        lbl = r.get("label", "unknown")
+        if lbl in ("ai", "human"):
+            source_label_counts[src][lbl] += 1
+
+    sources = sorted(source_label_counts.keys(),
+                     key=lambda s: sum(source_label_counts[s].values()),
+                     reverse=True)
+    ai_counts = [source_label_counts[s]["ai"] for s in sources]
+    human_counts = [source_label_counts[s]["human"] for s in sources]
+
+    fig, ax = plt.subplots(figsize=(10, 6))
+    x = np.arange(len(sources))
+    width = 0.38
+
+    b1 = ax.bar(x - width/2, human_counts, width,
+                label="Human", color=AURIS_BLUE, edgecolor="black", linewidth=0.3)
+    b2 = ax.bar(x + width/2, ai_counts, width,
+                label="AI", color=AURIS_RED, edgecolor="black", linewidth=0.3)
+
+    for bars in (b1, b2):
+        for bar in bars:
+            h = bar.get_height()
+            if h > 0:
+                ax.text(bar.get_x() + bar.get_width() / 2, h + 20,
+                        f"{int(h)}", ha="center", va="bottom", fontsize=8)
+
+    ax.set_xticks(x)
+    ax.set_xticklabels(sources, rotation=20, ha="right")
+    ax.set_ylabel("Samples")
+    ax.set_title("Class Balance per Source")
+    ax.legend(loc="upper right")
+
+    fig.tight_layout()
+    _save_fig(fig, output_dir, "fig10_per_source_balance")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Figure 11: t-SNE / PCA Feature Embedding
+# ═══════════════════════════════════════════════════════════════════════
+
+def plot_feature_embedding(
+    features_csv: Path,
+    output_dir: Path,
+    max_samples: int = 2000,
+) -> None:
+    """2D embedding (PCA + t-SNE) of features colored by class."""
+    try:
+        from sklearn.decomposition import PCA
+        from sklearn.manifold import TSNE
+        from sklearn.preprocessing import StandardScaler
+    except ImportError:
+        print("  Skipping embedding — sklearn not available")
+        return
+
+    X, y, _ = _load_features_with_names(features_csv)
+
+    if len(X) > max_samples:
+        rng = np.random.default_rng(42)
+        idx = rng.choice(len(X), max_samples, replace=False)
+        X = X[idx]
+        y = y[idx]
+
+    X_scaled = StandardScaler().fit_transform(X)
+
+    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+
+    # PCA
+    pca = PCA(n_components=2, random_state=42)
+    X_pca = pca.fit_transform(X_scaled)
+    ax = axes[0]
+    for cls, color, label in [(0, AURIS_BLUE, "Human"), (1, AURIS_RED, "AI")]:
+        mask = y == cls
+        ax.scatter(X_pca[mask, 0], X_pca[mask, 1],
+                   c=color, s=12, alpha=0.5, edgecolors="none", label=label)
+    var_exp = pca.explained_variance_ratio_
+    ax.set_xlabel(f"PC1 ({var_exp[0] * 100:.1f}%)")
+    ax.set_ylabel(f"PC2 ({var_exp[1] * 100:.1f}%)")
+    ax.set_title(f"PCA Projection (total var = {sum(var_exp) * 100:.1f}%)")
+    ax.legend(loc="best")
+
+    # t-SNE
+    try:
+        perplexity = min(30, max(5, len(X) // 20))
+        tsne = TSNE(n_components=2, perplexity=perplexity,
+                    random_state=42, max_iter=500, init="pca")
+        X_tsne = tsne.fit_transform(X_scaled)
+        ax = axes[1]
+        for cls, color, label in [(0, AURIS_BLUE, "Human"), (1, AURIS_RED, "AI")]:
+            mask = y == cls
+            ax.scatter(X_tsne[mask, 0], X_tsne[mask, 1],
+                       c=color, s=12, alpha=0.5, edgecolors="none", label=label)
+        ax.set_xlabel("t-SNE 1")
+        ax.set_ylabel("t-SNE 2")
+        ax.set_title(f"t-SNE Projection (perplexity={perplexity})")
+        ax.legend(loc="best")
+    except Exception as e:
+        print(f"  t-SNE failed: {e}")
+        axes[1].set_visible(False)
+
+    fig.suptitle("Feature Space Visualization — AI vs Human", fontsize=14, y=1.01)
+    fig.tight_layout()
+    _save_fig(fig, output_dir, "fig11_feature_embedding")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Figure 12: Probability Calibration
+# ═══════════════════════════════════════════════════════════════════════
+
+def plot_calibration_curves(
+    results: dict[str, Any],
+    output_dir: Path,
+    n_bins: int = 10,
+) -> None:
+    """Reliability diagram: predicted probability vs observed frequency."""
+    try:
+        from sklearn.calibration import calibration_curve
+    except ImportError:
+        return
+
+    fig, ax = plt.subplots(figsize=(7, 6))
+
+    for name, data in results.items():
+        if name.startswith("_"):
+            continue
+        y_true = np.array(data.get("y_true", []))
+        y_prob = np.array(data.get("y_prob", []))
+        if len(y_true) == 0:
+            continue
+        try:
+            frac_pos, mean_pred = calibration_curve(y_true, y_prob, n_bins=n_bins)
+            ax.plot(mean_pred, frac_pos, "o-",
+                    color=_get_color(name), label=name, linewidth=1.5, markersize=5)
+        except Exception:
+            continue
+
+    ax.plot([0, 1], [0, 1], "k--", alpha=0.5, label="Perfect calibration")
+    ax.set_xlabel("Mean Predicted Probability")
+    ax.set_ylabel("Fraction of Positives")
+    ax.set_title("Probability Calibration (Reliability Diagram)")
+    ax.legend(loc="lower right", fontsize=8)
+    ax.set_xlim([0, 1])
+    ax.set_ylim([0, 1])
+    ax.set_aspect("equal")
+
+    fig.tight_layout()
+    _save_fig(fig, output_dir, "fig12_calibration")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Figure 13: Prediction Score Distribution
+# ═══════════════════════════════════════════════════════════════════════
+
+def plot_score_distribution(
+    results: dict[str, Any],
+    output_dir: Path,
+) -> None:
+    """Best model's predicted probability distribution per class."""
+    best_model = results.get("_best_model", "")
+    if not best_model or best_model not in results:
+        candidates = [k for k in results if not k.startswith("_")]
+        if not candidates:
+            return
+        best_model = candidates[0]
+
+    data = results[best_model]
+    y_true = np.array(data.get("y_true", []))
+    y_prob = np.array(data.get("y_prob", []))
+    if len(y_true) == 0:
+        return
+
+    fig, ax = plt.subplots(figsize=(9, 5))
+
+    ax.hist(y_prob[y_true == 0], bins=40, alpha=0.6,
+            color=AURIS_BLUE, label="Human (true class)", edgecolor="white")
+    ax.hist(y_prob[y_true == 1], bins=40, alpha=0.6,
+            color=AURIS_RED, label="AI (true class)", edgecolor="white")
+    ax.axvline(0.5, color="black", linestyle="--", linewidth=1, alpha=0.7,
+               label="Decision threshold (0.5)")
+
+    ax.set_xlabel(f"Predicted P(AI) — {best_model}")
+    ax.set_ylabel("Count")
+    ax.set_title("Prediction Score Distribution — Best Model")
+    ax.legend(loc="upper center")
+
+    fig.tight_layout()
+    _save_fig(fig, output_dir, "fig13_score_distribution")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Extended Summary Statistics Table (Markdown)
+# ═══════════════════════════════════════════════════════════════════════
+
+def generate_extended_summary(
+    results: dict[str, Any],
+    manifest_csv: Path,
+    features_csv: Path,
+    output_dir: Path,
+) -> None:
+    """Comprehensive markdown summary for the academic report."""
+    rows = _load_manifest(manifest_csv)
+    X, y, feature_cols = _load_features_with_names(features_csv)
+
+    from collections import Counter
+    labels = Counter(r.get("label", "?") for r in rows)
+    sources = Counter(r.get("generator", "?") for r in rows)
+
+    model_names = [k for k in results if not k.startswith("_")]
+    best_model = results.get("_best_model", "")
+
+    lines = []
+    lines.append("# AURIS — Dataset & Results Summary\n")
+
+    lines.append("## Dataset\n")
+    lines.append(f"- **Total samples:** {len(rows):,}")
+    lines.append(f"- **Features extracted:** {len(feature_cols)}")
+    lines.append(f"- **Samples with extracted features:** {len(X):,}")
+    lines.append("")
+    lines.append("### Class distribution")
+    lines.append("| Class | Count | % |")
+    lines.append("|-------|-------|---|")
+    total = sum(labels.values())
+    for lbl in sorted(labels):
+        pct = labels[lbl] / total * 100 if total else 0
+        lines.append(f"| {lbl} | {labels[lbl]:,} | {pct:.1f}% |")
+    lines.append("")
+
+    lines.append("### Source distribution")
+    lines.append("| Source | Count | % |")
+    lines.append("|--------|-------|---|")
+    for src in sorted(sources, key=lambda s: sources[s], reverse=True):
+        pct = sources[src] / total * 100 if total else 0
+        lines.append(f"| {src} | {sources[src]:,} | {pct:.1f}% |")
+    lines.append("")
+
+    lines.append("## Results\n")
+    lines.append(f"- **Best model:** {best_model}")
+    n_folds = results.get("_n_folds", "?")
+    lines.append(f"- **Cross-validation folds:** {n_folds}")
+    lines.append("")
+
+    lines.append("### Model performance")
+    lines.append("| Model | Accuracy | Precision | Recall | F1 | ROC-AUC |")
+    lines.append("|-------|----------|-----------|--------|-----|---------|")
+    for name in model_names:
+        d = results[name]
+        bold = "**" if name == best_model else ""
+        lines.append(
+            f"| {bold}{name}{bold} | "
+            f"{d.get('accuracy', 0):.4f} | "
+            f"{d.get('precision', 0):.4f} | "
+            f"{d.get('recall', 0):.4f} | "
+            f"{d.get('f1', 0):.4f} | "
+            f"{d.get('roc_auc', 0):.4f} |"
+        )
+    lines.append("")
+
+    if "_feature_importance" in results:
+        lines.append("### Top-15 Feature Importances")
+        lines.append("| Rank | Feature | Importance |")
+        lines.append("|------|---------|------------|")
+        sorted_imp = sorted(
+            results["_feature_importance"].items(),
+            key=lambda x: x[1], reverse=True,
+        )[:15]
+        for i, (feat, imp) in enumerate(sorted_imp, 1):
+            lines.append(f"| {i} | {feat} | {imp:.4f} |")
+        lines.append("")
+
+    output_path = output_dir / "REPORT_SUMMARY.md"
+    output_path.write_text("\n".join(lines), encoding="utf-8")
+    print(f"  Saved: REPORT_SUMMARY.md")
+
+
 # ═══════════════════════════════════════════════════════════════════════
 # Main entry point
 # ═══════════════════════════════════════════════════════════════════════