plotfeatures.py

# plotfeatures.py
# -------------------------------------------------------------
# Generates 2D scatter feature plots from a combined feature CSV.
# Uses label inferred from event name (event_1_, event_2_, etc.)
# Saves all plots as JPEG files in feature_plots/.
# -------------------------------------------------------------

import os
from datetime import datetime

import pandas as pd
import matplotlib.pyplot as plt
from tkinter import Tk, filedialog

# 0 will mean "Unlabeled" (event_0_...)
LABEL_NAMES = {
    0: "Unlabeled",
    1: "Collision",
    2: "Drop",
    3: "Stoppage"
}

COLORS = {
    0: "tab:brown",
    1: "tab:blue",
    2: "tab:orange",
    3: "tab:green"
}

MARKERS = {
    0: "P",
    1: "o",
    2: "s",
    3: "^"
}


def get_label_from_event(name):
    parts = name.split("_")
    if len(parts) > 1:
        try:
            return int(parts[1])
        except:
            return 0
    return 0


def main():
    # Select feature CSV
    print("[Message]: Select combined feature CSV...")
    root = Tk()
    root.withdraw()
    path = filedialog.askopenfilename(
        initialdir="event_features",
        title="Select combined feature CSV",
        filetypes=[("CSV files", "*.csv")]
    )
    root.destroy()

    if not path:
        return

    print("[Message]: Loading CSV...")
    df = pd.read_csv(path)

    # Add label column if needed
    if "label" not in df.columns and "event" in df.columns:
        df["label"] = df["event"].apply(get_label_from_event)

    # Collect numeric feature columns
    numeric_cols = [
        c for c in df.columns
        if c.endswith(("_mean", "_std", "_rms", "_ptp", "_min", "_max"))
    ]

    # Build all pairs
    feature_pairs = []
    for i in range(len(numeric_cols)):
        for j in range(i + 1, len(numeric_cols)):
            feature_pairs.append((numeric_cols[i], numeric_cols[j]))

    # Output folder
    ts = datetime.now().strftime("%m%d%Y_%H%M%S")
    base_dir = "feature_plots"
    out_dir = os.path.join(base_dir, f"plots_{ts}")
    os.makedirs(out_dir, exist_ok=True)

    print("[Message]: Generating plots...")
    print(f"[Message]: Total plots = {len(feature_pairs)}")

    # Generate scatter plots
    for x, y in feature_pairs:
        plt.figure(figsize=(7, 6))

        for lbl in sorted(df["label"].unique()):
            sub = df[df["label"] == lbl]
            plt.scatter(
                sub[x],
                sub[y],
                c=COLORS.get(lbl, "k"),
                marker=MARKERS.get(lbl, "o"),
                label=LABEL_NAMES.get(lbl, f"Label {lbl}"),
                alpha=0.8
            )

        plt.xlabel(x)
        plt.ylabel(y)
        plt.title(f"{x} vs {y}")
        plt.grid(True)
        plt.legend()
        plt.tight_layout()

        filename = f"{x}_vs_{y}.jpeg"
        plt.savefig(os.path.join(out_dir, filename), dpi=300)
        plt.close()

    print("[Saved]:", out_dir)
    print("[Message]: All plots generated.\n")


if __name__ == "__main__":
    main()