scripts/eda.py

#!/usr/bin/env python3
"""Exploratory Data Analysis for NIH Chest X-ray14 dataset.

Streams metadata from HuggingFace (no heavy downloads) and generates
summary statistics and plots.

Usage:
    python scripts/eda.py
    python scripts/eda.py --num-samples 10000 --output-dir results/eda
"""

from __future__ import annotations

import argparse
from collections import Counter
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from datasets import load_dataset

# Standalone copy — no import from src so this script is self-contained.
PATHOLOGY_LABELS = [
    "Atelectasis",
    "Cardiomegaly",
    "Effusion",
    "Infiltration",
    "Mass",
    "Nodule",
    "Pneumonia",
    "Pneumothorax",
    "Consolidation",
    "Edema",
    "Emphysema",
    "Fibrosis",
    "Pleural_Thickening",
    "Hernia",
]

DATASET_REPO = "alkzar90/NIH-Chest-X-ray-dataset"


def stream_samples(num_samples: int) -> list[dict]:
    """Stream *num_samples* rows from the HF dataset (metadata + images)."""
    print(f"Streaming {num_samples:,} samples from {DATASET_REPO} ...")
    ds = load_dataset(DATASET_REPO, streaming=True)
    samples = list(ds["train"].take(num_samples))
    print(f"Loaded {len(samples):,} samples.")
    return samples


# ---------------------------------------------------------------------------
# Plot helpers
# ---------------------------------------------------------------------------

def plot_label_distribution(df: pd.DataFrame, output_dir: Path) -> None:
    """Horizontal bar chart of label counts."""
    label_counts: Counter[str] = Counter()
    for labels_str in df["labels"]:
        if pd.isna(labels_str) or labels_str == "No Finding":
            label_counts["No Finding"] += 1
        else:
            for label in str(labels_str).split("|"):
                label_counts[label.strip()] += 1

    sorted_items = sorted(label_counts.items(), key=lambda x: -x[1])
    labels, counts = zip(*sorted_items)

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.barh(labels, counts, color="steelblue")
    ax.set_xlabel("Count")
    ax.set_title("Label Distribution — NIH Chest X-ray14 (EDA sample)")
    plt.tight_layout()
    fig.savefig(output_dir / "label_distribution.png", dpi=150)
    plt.close(fig)
    print("  Saved label_distribution.png")


def plot_binary_distribution(df: pd.DataFrame, output_dir: Path) -> None:
    """Pie chart of Normal vs Abnormal."""
    is_abnormal = df["labels"].apply(
        lambda x: 0 if pd.isna(x) or x == "No Finding" else 1
    )
    normal_count = int((is_abnormal == 0).sum())
    abnormal_count = int((is_abnormal == 1).sum())

    fig, ax = plt.subplots(figsize=(6, 5))
    ax.pie(
        [normal_count, abnormal_count],
        labels=["Normal", "Abnormal"],
        autopct="%1.1f%%",
        colors=["#66b3ff", "#ff6666"],
        startangle=90,
    )
    ax.set_title("Normal vs Abnormal Distribution")
    plt.tight_layout()
    fig.savefig(output_dir / "binary_distribution.png", dpi=150)
    plt.close(fig)
    print("  Saved binary_distribution.png")


def plot_cooccurrence_matrix(df: pd.DataFrame, output_dir: Path) -> None:
    """Heatmap of pathology co-occurrences."""
    n = len(PATHOLOGY_LABELS)
    cooccurrence = np.zeros((n, n), dtype=int)

    for labels_str in df["labels"]:
        if pd.isna(labels_str) or labels_str == "No Finding":
            continue
        present = [
            PATHOLOGY_LABELS.index(lbl.strip())
            for lbl in str(labels_str).split("|")
            if lbl.strip() in PATHOLOGY_LABELS
        ]
        for i in present:
            for j in present:
                cooccurrence[i][j] += 1

    fig, ax = plt.subplots(figsize=(12, 10))
    sns.heatmap(
        cooccurrence,
        xticklabels=PATHOLOGY_LABELS,
        yticklabels=PATHOLOGY_LABELS,
        cmap="YlOrRd",
        annot=True,
        fmt="d",
        ax=ax,
    )
    ax.set_title("Pathology Co-occurrence Matrix")
    plt.tight_layout()
    fig.savefig(output_dir / "cooccurrence_matrix.png", dpi=150)
    plt.close(fig)
    print("  Saved cooccurrence_matrix.png")


def plot_sample_images(samples: list[dict], output_dir: Path) -> None:
    """2x5 grid of sample X-ray images (streamed from HF)."""
    fig, axes = plt.subplots(2, 5, figsize=(20, 8))
    for ax, sample in zip(axes.flat, samples[:10]):
        img = sample["image"]
        ax.imshow(img, cmap="gray")
        label_text = str(sample.get("labels", sample.get("label", "?")))[:30]
        ax.set_title(label_text, fontsize=8)
        ax.axis("off")

    fig.suptitle("Sample Chest X-ray Images (streamed from HuggingFace)", fontsize=14)
    plt.tight_layout()
    fig.savefig(output_dir / "sample_images.png", dpi=150)
    plt.close(fig)
    print("  Saved sample_images.png")


# ---------------------------------------------------------------------------
# Summary statistics
# ---------------------------------------------------------------------------

def print_summary(df: pd.DataFrame) -> None:
    """Print key statistics to stdout."""
    total = len(df)
    is_abnormal = df["labels"].apply(
        lambda x: 0 if pd.isna(x) or x == "No Finding" else 1
    )
    normal_count = int((is_abnormal == 0).sum())
    abnormal_count = int((is_abnormal == 1).sum())

    label_counts: Counter[str] = Counter()
    multi_label_count = 0
    for labels_str in df["labels"]:
        if pd.isna(labels_str) or labels_str == "No Finding":
            continue
        individual = [lbl.strip() for lbl in str(labels_str).split("|")]
        if len(individual) > 1:
            multi_label_count += 1
        for label in individual:
            label_counts[label] += 1

    print("\n" + "=" * 60)
    print("  CheXVision — EDA Summary Statistics")
    print("=" * 60)
    print(f"  Total samples analysed : {total:,}")
    print(f"  Normal (No Finding)    : {normal_count:,} ({normal_count / total:.1%})")
    print(f"  Abnormal               : {abnormal_count:,} ({abnormal_count / total:.1%})")
    print(f"  Multi-label samples    : {multi_label_count:,} ({multi_label_count / total:.1%})")
    print("-" * 60)
    print("  Per-pathology counts:")
    for label in PATHOLOGY_LABELS:
        cnt = label_counts.get(label, 0)
        print(f"    {label:<22s} {cnt:>6,}  ({cnt / total:.2%})")
    print("=" * 60 + "\n")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Lightweight EDA for NIH Chest X-ray14 (streams from HuggingFace)."
    )
    parser.add_argument(
        "--num-samples",
        type=int,
        default=5000,
        help="Number of samples to stream (default: 5000).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("results/eda"),
        help="Directory for output plots (default: results/eda).",
    )
    args = parser.parse_args()

    sns.set_theme(style="whitegrid")
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # Stream samples
    samples = stream_samples(args.num_samples)

    # Build a lightweight DataFrame (labels only — no images)
    records = []
    for s in samples:
        labels = s.get("labels", s.get("label", "No Finding"))
        records.append({"labels": labels})
    df = pd.DataFrame(records)

    # Generate plots
    print("Generating plots ...")
    plot_label_distribution(df, args.output_dir)
    plot_binary_distribution(df, args.output_dir)
    plot_cooccurrence_matrix(df, args.output_dir)
    plot_sample_images(samples, args.output_dir)

    # Print summary
    print_summary(df)

    print(f"All plots saved to {args.output_dir.resolve()}")


if __name__ == "__main__":
    main()