my_project.plots

  1"""
  2Utilities to evaluate a trained model and generate performance visualizations.
  3- evaluate_and_plot(...) : runs eval on a dataloader and saves:
  4    * confusion_matrix.png
  5    * per_class_accuracy.png
  6    * misclassified_grid.png
  7- plot_curves_from_csvlogger(...) : plots curves from Lightning's CSVLogger
  8    * train_loss_[step|epoch].png
  9    * val_acc_epoch.png (if logged during validation)
 10Outputs are saved under reports/figures/ by default.
 11
 12This module does NOT modify your model architecture.
 13"""
 14
 15import os
 16import math
 17from typing import Tuple, List, Optional
 18
 19import torch
 20import torch.nn.functional as F
 21import matplotlib.pyplot as plt
 22import numpy as np
 23import pandas as pd
 24from sklearn.metrics import confusion_matrix
 25from sklearn.calibration import calibration_curve
 26from scipy.cluster import hierarchy as sch
 27from scipy.spatial.distance import pdist
 28
 29
 30# Fashion-MNIST class names (0..9)
 31FASHION_CLASSES = [
 32    "T-shirt/top",
 33    "Trouser",
 34    "Pullover",
 35    "Dress",
 36    "Coat",
 37    "Sandal",
 38    "Shirt",
 39    "Sneaker",
 40    "Bag",
 41    "Ankle boot",
 42]
 43
 44
 45def _ensure_dir(path: str) -> None:
 46    os.makedirs(path, exist_ok=True)
 47
 48
 49def _to_numpy(x: torch.Tensor) -> np.ndarray:
 50    return x.detach().cpu().numpy()
 51
 52
 53def _get_device(module: torch.nn.Module) -> torch.device:
 54    return next(module.parameters()).device
 55
 56
 57@torch.no_grad()
 58def evaluate_model(
 59    model: torch.nn.Module,
 60    dataloader: torch.utils.data.DataLoader,
 61) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
 62    """
 63    Run the model over a dataloader and return predictions.
 64
 65    Parameters
 66    ----------
 67    model : torch.nn.Module
 68        Trained model to evaluate.
 69    dataloader : torch.utils.data.DataLoader
 70        DataLoader for evaluation (test/validation).
 71
 72    Returns
 73    -------
 74    tuple
 75        (y_true, y_pred, y_prob_max):
 76        - y_true (ndarray[int]): Ground-truth labels.
 77        - y_pred (ndarray[int]): Predicted labels.
 78        - y_prob_max (ndarray[float]): Max softmax confidence per sample.
 79    """
 80
 81    model.eval()
 82    device = _get_device(model)
 83
 84    all_true: List[int] = []
 85    all_pred: List[int] = []
 86    all_conf: List[float] = []
 87
 88    for xb, yb in dataloader:
 89        xb = xb.to(device)
 90        yb = yb.to(device)
 91
 92        logits = model(xb)
 93        probs = F.softmax(logits, dim=1)
 94        conf, pred = probs.max(dim=1)
 95
 96        all_true.extend(_to_numpy(yb))
 97        all_pred.extend(_to_numpy(pred))
 98        all_conf.extend(_to_numpy(conf))
 99
100    return (
101        np.asarray(all_true, dtype=np.int64),
102        np.asarray(all_pred, dtype=np.int64),
103        np.asarray(all_conf, dtype=np.float32),
104    )
105
106
107def plot_confusion_matrix(
108    y_true: np.ndarray,
109    y_pred: np.ndarray,
110    class_names: List[str],
111    out_path: str,
112    normalize: bool = True,
113) -> None:
114    """
115    Plot and save a confusion matrix.
116
117    Parameters
118    ----------
119    y_true : np.ndarray
120        Ground-truth labels.
121    y_pred : np.ndarray
122        Predicted labels.
123    class_names : list of str
124        Class names corresponding to label indices.
125    out_path : str
126        File path to save the plot.
127    normalize : bool, optional (default=True)
128        If True, normalize counts to percentages.
129    """
130
131    cm = confusion_matrix(y_true, y_pred, labels=list(range(len(class_names))))
132    if normalize:
133        with np.errstate(all="ignore"):
134            cm = cm.astype(np.float64) / cm.sum(axis=1, keepdims=True)
135        cm = np.nan_to_num(cm)
136
137    plt.figure(figsize=(6, 5))
138    plt.imshow(cm, interpolation="nearest", cmap="coolwarm")
139    plt.title("Confusion Matrix " if normalize else "Confusion Matrix")
140    plt.colorbar()
141    tick_marks = np.arange(len(class_names))
142    plt.xticks(tick_marks, class_names, rotation=45, ha="right")
143    plt.yticks(tick_marks, class_names)
144
145    # Annotate cells
146    thresh = cm.max() / 2.0 if cm.size > 0 else 0.5
147    for i in range(cm.shape[0]):
148        for j in range(cm.shape[1]):
149            val = cm[i, j]
150            txt = f"{val:.2f}" if normalize else f"{int(val)}"
151            plt.text(
152                j,
153                i,
154                txt,
155                horizontalalignment="center",
156                verticalalignment="center",
157                fontsize=8,
158                color="white" if cm[i, j] > thresh else "black",
159            )
160
161    plt.ylabel("True label")
162    plt.xlabel("Predicted label")
163    plt.tight_layout()
164    plt.savefig(out_path, dpi=160)
165    plt.close()
166
167
168def plot_per_class_accuracy(
169    y_true: np.ndarray,
170    y_pred: np.ndarray,
171    class_names: List[str],
172    out_path: str,
173) -> None:
174    """
175    Plot and save per-class accuracy.
176
177    Parameters
178    ----------
179    y_true : np.ndarray
180        Ground-truth labels.
181    y_pred : np.ndarray
182        Predicted labels.
183    class_names : list of str
184        Class names corresponding to label indices.
185    out_path : str
186        File path to save the plot.
187    """
188
189    num_classes = len(class_names)
190    correct = np.zeros(num_classes, dtype=np.int64)
191    total = np.zeros(num_classes, dtype=np.int64)
192
193    for t, p in zip(y_true, y_pred):
194        total[t] += 1
195        if t == p:
196            correct[t] += 1
197
198    acc = np.divide(correct, np.maximum(total, 1), where=total > 0)
199
200    plt.figure(figsize=(9, 4))
201    plt.bar(np.arange(num_classes), acc)
202    plt.ylim(0, 1)
203    plt.xticks(np.arange(num_classes), class_names, rotation=45, ha="right")
204    plt.ylabel("Accuracy")
205    plt.title("Per-class Accuracy")
206    plt.tight_layout()
207    plt.savefig(out_path, dpi=160)
208    plt.close()
209
210
211def _denormalize_img(img: torch.Tensor) -> np.ndarray:
212    """
213    Convert a normalized tensor image back to NumPy format.
214
215    Parameters
216    ----------
217    img : torch.Tensor
218        Image tensor (CHW or HW).
219
220    Returns
221    -------
222    np.ndarray
223        Image array suitable for plotting.
224    """
225
226    if img.dim() == 3 and img.shape[0] == 1:
227        img_np = img.squeeze(0).detach().cpu().numpy()
228        return img_np
229    # If already HxW
230    return img.detach().cpu().numpy()
231
232
233@torch.no_grad()
234def plot_misclassified_grid(
235    model: torch.nn.Module,
236    dataloader: torch.utils.data.DataLoader,
237    class_names: List[str],
238    out_path: str,
239    max_examples: int = 16,
240) -> None:
241    """
242    Plot a grid of misclassified samples.
243
244    Parameters
245    ----------
246    model : torch.nn.Module
247        Trained model to evaluate.
248    dataloader : torch.utils.data.DataLoader
249        DataLoader for evaluation (test/validation).
250    class_names : list of str
251        Class names corresponding to label indices.
252    out_path : str
253        File path to save the plot.
254    max_examples : int, optional (default=16)
255        Maximum number of misclassified samples to show.
256    """
257    model.eval()
258    device = _get_device(model)
259
260    images = []
261    labels_true = []
262    labels_pred = []
263
264    # Collect up to max_examples misclassified samples
265    for xb, yb in dataloader:
266        xb = xb.to(device)
267        yb = yb.to(device)
268        logits = model(xb)
269        pred = logits.argmax(dim=1)
270
271        mism = pred != yb
272        if mism.any():
273            mis_idx = torch.nonzero(mism).flatten()
274            for idx in mis_idx:
275                images.append(xb[idx].cpu())
276                labels_true.append(int(yb[idx].cpu()))
277                labels_pred.append(int(pred[idx].cpu()))
278                if len(images) >= max_examples:
279                    break
280        if len(images) >= max_examples:
281            break
282
283    if len(images) == 0:
284        # Nothing to plot
285        fig = plt.figure(figsize=(6, 2))
286        plt.text(0.5, 0.5, "No misclassified samples found.", ha="center", va="center")
287        plt.axis("off")
288        plt.tight_layout()
289        fig.savefig(out_path, dpi=160)
290        plt.close()
291        return
292
293    # Determine grid size
294    cols = int(math.ceil(math.sqrt(len(images))))
295    rows = int(math.ceil(len(images) / cols))
296
297    plt.figure(figsize=(cols * 1.5, rows * 1.5))
298    for i, img in enumerate(images):
299        ax = plt.subplot(rows, cols, i + 1)
300        ax.imshow(_denormalize_img(img), cmap="gray")
301        true_name = class_names[labels_true[i]]
302        pred_name = class_names[labels_pred[i]]
303        ax.set_title(f"{true_name} → {pred_name}", fontsize=8)
304        ax.axis("off")
305
306    plt.suptitle("Misclassified Samples", y=0.98)
307    plt.tight_layout()
308    plt.savefig(out_path, dpi=160)
309    plt.close()
310
311
312def evaluate_and_plot(
313    model: torch.nn.Module,
314    datamodule,
315    out_dir: str = "reports/figures",
316) -> dict:
317
318    _ensure_dir(out_dir)
319
320    # Use the datamodule's test dataloader (Lightning-style)
321    test_loader = datamodule.test_dataloader()
322
323    # 1) Predictions + basic accuracy
324    y_true, y_pred, y_conf = evaluate_model(model, test_loader)
325    test_acc = float((y_true == y_pred).mean())
326
327    # 2) Confusion matrix (normalized)
328    cm_path = os.path.join(out_dir, "confusion_matrix.png")
329    plot_confusion_matrix(y_true, y_pred, FASHION_CLASSES, cm_path, normalize=True)
330
331    # 3) Per-class accuracy bars
332    pca_path = os.path.join(out_dir, "per_class_accuracy.png")
333    plot_per_class_accuracy(y_true, y_pred, FASHION_CLASSES, pca_path)
334
335    # 4) Misclassified samples grid
336    mis_path = os.path.join(out_dir, "misclassified_grid.png")
337    plot_misclassified_grid(
338        model, test_loader, FASHION_CLASSES, mis_path, max_examples=6
339    )
340
341    # 5) Calibration (reliability) curve
342    calib_path = os.path.join(out_dir, "calibration_curve.png")
343    y_correct = (y_true == y_pred).astype(int)
344    plot_calibration_curve(y_correct, y_conf, calib_path, n_bins=10)
345
346    return {
347        "test_accuracy": test_acc,
348        "confusion_matrix": cm_path,
349        "per_class_accuracy": pca_path,
350        "misclassified_grid": mis_path,
351        "calibration_curve": calib_path,
352    }
353
354
355def plot_calibration_curve(y_true, y_prob, out_path, n_bins: int = 10):
356
357    # Calibration curve
358    prob_true, prob_pred = calibration_curve(
359        y_true, y_prob, n_bins=n_bins, strategy="uniform"
360    )
361
362    plt.figure(figsize=(6, 5))
363    plt.plot(prob_pred, prob_true, marker="o", label="Model")
364    plt.plot([0, 1], [0, 1], linestyle="--", color="gray", label="Perfectly calibrated")
365    plt.xlabel("Predicted probability")
366    plt.ylabel("Observed frequency")
367    plt.title("Calibration Curve")
368    plt.legend()
369    plt.tight_layout()
370    plt.savefig(out_path, dpi=160)
371    plt.close()
372
373
374def plot_curves_from_csvlogger(
375    csv_log_dir: str,
376    out_dir: str = "reports/figures",
377    train_loss_keys: Optional[List[str]] = None,
378    val_acc_keys: Optional[List[str]] = None,
379) -> List[str]:
380    """
381    If you use Lightning's CSVLogger, this parses the `metrics.csv` file and
382    produces line plots for:
383      - train loss (step/epoch)
384      - val accuracy (epoch)
385
386    Usage in train.py:
387        from pytorch_lightning.loggers import CSVLogger
388        logger = CSVLogger("logs", name="fashion")
389        trainer = pl.Trainer(..., logger=logger)
390
391    Then call:
392        plot_curves_from_csvlogger(logger.log_dir)
393
394    Args:
395        csv_log_dir: e.g., "logs/fashion/version_0"
396        train_loss_keys: possible metric column names for train loss
397        val_acc_keys: possible metric column names for val acc
398
399    Returns:
400        list of saved figure paths
401    """
402    _ensure_dir(out_dir)
403    metrics_path = os.path.join(csv_log_dir, "metrics.csv")
404    if not os.path.exists(metrics_path):
405        return []
406
407    df = pd.read_csv(metrics_path)
408
409    # Reasonable defaults; adjust if you log different names
410    if train_loss_keys is None:
411        # Lightning often stores as "train_loss_step" and/or "train_loss_epoch"
412        train_loss_keys = ["train_loss_step", "train_loss_epoch", "loss", "train_loss"]
413    if val_acc_keys is None:
414        # You already log 'val_acc' in your model
415        val_acc_keys = ["val_acc", "val_acc_epoch"]
416
417    saved = []
418
419    # Train loss vs step/epoch
420    for key in train_loss_keys:
421        if key in df.columns:
422            # Prefer 'step' index if present, else 'epoch'
423            if "step" in df.columns and not df["step"].isna().all():
424                x = df["step"]
425                x_label = "Step"
426            else:
427                x = df["epoch"] if "epoch" in df.columns else np.arange(len(df))
428                x_label = "Epoch"
429
430            plt.figure(figsize=(6, 4))
431            plt.plot(x, df[key])
432            plt.xlabel(x_label)
433            plt.ylabel("Train Loss")
434            plt.title(f"{key} over {x_label.lower()}")
435            plt.tight_layout()
436            out_path = os.path.join(out_dir, f"{key}.png")
437            plt.savefig(out_path, dpi=160)
438            plt.close()
439            saved.append(out_path)
440            # Only plot first that exists (avoid duplicates)
441            break
442
443    # Validation accuracy vs epoch
444    for key in val_acc_keys:
445        if key in df.columns:
446            x = df["epoch"] if "epoch" in df.columns else np.arange(len(df))
447            plt.figure(figsize=(6, 4))
448            plt.plot(x, df[key])
449            plt.xlabel("Epoch")
450            plt.ylabel("Validation Accuracy")
451            plt.title(f"{key} over epochs")
452            plt.tight_layout()
453            out_path = os.path.join(out_dir, f"{key}.png")
454            plt.savefig(out_path, dpi=160)
455            plt.close()
456            saved.append(out_path)
457            break
458
459    return saved
460
461
462def plot_learning_curves_from_df(
463    df: pd.DataFrame,
464) -> Tuple[Optional[plt.Figure], Optional[plt.Figure]]:
465    """
466    Parses a metrics DataFrame and produces line plots for train loss and val accuracy.
467
468    Args:
469        df (pd.DataFrame): DataFrame with metrics (e.g., from CSVLogger).
470
471    Returns:
472        A tuple of (train_loss_figure, val_acc_figure). Figures can be None.
473    """
474    train_loss_fig, val_acc_fig = None, None
475
476    # Train loss vs step/epoch
477    train_loss_keys = ["train_loss_step", "train_loss_epoch", "loss", "train_loss"]
478    val_acc_keys = ["val_acc", "val_acc_epoch"]
479
480    for key in train_loss_keys:
481        if key in df.columns and not df[key].isna().all():
482            series = df[key].dropna()
483            # Prefer 'step' index if present, else 'epoch'
484            if "step" in df.columns and not df["step"].isna().all():
485                x = df.loc[series.index, "step"]
486                x_label = "Step"
487            else:
488                x = df.loc[series.index, "epoch"]
489                x_label = "Epoch"
490
491            train_loss_fig = plt.figure(figsize=(6, 4))
492            plt.plot(x, series)
493            plt.xlabel(x_label)
494            plt.ylabel("Train Loss")
495            plt.title(f"Training Loss vs. {x_label}")
496            plt.tight_layout()
497            # Only plot first that exists
498            break
499
500    # Validation accuracy vs epoch
501    for key in val_acc_keys:
502        if key in df.columns and not df[key].isna().all():
503            # Filter to get only the steps where validation was run
504            series = df[key].dropna()
505            x = df.loc[series.index, "epoch"]
506            val_acc_fig = plt.figure(figsize=(6, 4))
507            plt.plot(x, series, marker='o', linestyle='-')
508            plt.xlabel("Epoch")
509            plt.ylabel("Validation Accuracy")
510            plt.title("Validation Accuracy vs. Epoch")
511            plt.ylim(0, 1)
512            plt.tight_layout()
513            # Only plot first that exists
514            break
515
516    # Close figures if they were not created to avoid empty plots
517    if train_loss_fig is None: plt.close()
518    if val_acc_fig is None: plt.close()
519
520    return train_loss_fig, val_acc_fig
521
522
523def plot_class_distribution(
524    df: pd.DataFrame, class_names: List[str], out_path: Optional[str] = None
525) -> plt.Figure:
526    """
527    Plots the distribution of classes in a dataset.
528
529    Args:
530        df (pd.DataFrame): DataFrame with a 'label' column.
531        class_names (List[str]): List of class names.
532        out_path (Optional[str], optional): If provided, saves the plot. Defaults to None.
533
534    Returns:
535        plt.Figure: The matplotlib figure object.
536    """
537    class_counts = df["label"].value_counts().sort_index()
538    total_samples = class_counts.sum()
539
540    fig = plt.figure(figsize=(9, 4))
541    bars = plt.bar(class_counts.index, class_counts.values)
542    plt.xticks(np.arange(len(class_names)), class_names, rotation=45, ha="right")
543    plt.ylabel("Number of Samples")
544    plt.title("Dataset Class Distribution")
545
546    # Add count and percentage labels
547    for bar in bars:
548        height = bar.get_height()
549        plt.text(bar.get_x() + bar.get_width()/2., height, f'{height}\n({height/total_samples:.1%})', ha='center', va='bottom', fontsize=8)
550
551    plt.tight_layout()
552
553    if out_path:
554        _ensure_dir(os.path.dirname(out_path))
555        plt.savefig(out_path, dpi=160)
556        plt.close()
557
558    return fig
559
560
561def get_sample_images_for_gallery(
562    df: pd.DataFrame, class_names: List[str], n_samples: int = 20
563) -> List[Tuple[np.ndarray, str]]:
564    """
565    Gets a list of random sample images and their labels for a Gradio Gallery.
566
567    Args:
568        df (pd.DataFrame): DataFrame with image data.
569        class_names (List[str]): List of class names.
570        n_samples (int, optional): Number of samples to retrieve. Defaults to 20.
571
572    Returns:
573        List[Tuple[np.ndarray, str]]: A list of tuples, each containing an image and its label.
574    """
575    samples = []
576    if len(df) > n_samples:
577        df_sample = df.sample(n=n_samples)
578    else:
579        df_sample = df
580
581    for _, row in df_sample.iterrows():
582        label = class_names[int(row["label"])]
583        image_data = row.iloc[1:].values.astype(np.uint8)
584        image = image_data.reshape(28, 28)
585        # The label is now just plain text
586        samples.append((image, label))
587    return samples
588
589
590def plot_class_correlation_dendrogram(
591    df: pd.DataFrame, class_names: List[str], out_path: Optional[str] = None
592) -> plt.Figure:
593    """
594    Calculates and plots a dendrogram showing the similarity between the
595    average image of each class.
596
597    Args:
598        df (pd.DataFrame): DataFrame with image data and labels.
599        class_names (List[str]): List of class names.
600        out_path (Optional[str], optional): If provided, saves the plot. Defaults to None.
601
602    Returns:
603        plt.Figure: The matplotlib figure object.
604    """
605    mean_images = []
606    for i in range(len(class_names)):
607        # Filter for the class, get pixel data, and calculate the mean image
608        mean_img = df[df["label"] == i].iloc[:, 1:].mean().values
609        mean_images.append(mean_img)
610
611    # Convert list of mean images to a matrix
612    mean_images_matrix = np.vstack(mean_images)
613
614    # Perform hierarchical clustering
615    linked = sch.linkage(mean_images_matrix, method="ward")
616
617    fig = plt.figure(figsize=(9, 4))
618    sch.dendrogram(linked, orientation="top", labels=class_names, leaf_rotation=90)
619    plt.ylabel("Euclidean Distance (between mean images)")
620    plt.title("Class Similarity Dendrogram")
621    plt.tight_layout()
622
623    if out_path:
624        _ensure_dir(os.path.dirname(out_path))
625        plt.savefig(out_path, dpi=160)
626        plt.close()
627
628    return fig