92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from sklearn.metrics import classification_report
|
|
import matplotlib
|
|
from matplotlib import pyplot as plt
|
|
|
|
from scripts.training.mobilenetv3 import mobilenetv3, RFClassifier
|
|
from helpers.app_settings import get_app_settings
|
|
from cm_plotter import plot_confusion_matrix
|
|
from scripts.training.modulation_dataset import ModulationH5Dataset
|
|
|
|
|
|
def load_validation_data():
|
|
val_dataset = ModulationH5Dataset(
|
|
"data/dataset/val.h5",
|
|
label_name="modulation",
|
|
data_key="validation_data"
|
|
)
|
|
|
|
X = np.stack([x.numpy() for x, _ in val_dataset]) # shape: (N, C, L)
|
|
y = np.array([y.item() for _, y in val_dataset]) # shape: (N,)
|
|
class_names = list(val_dataset.label_encoder.classes_)
|
|
|
|
return X, y, class_names
|
|
|
|
|
|
def build_model_from_ckpt(
|
|
ckpt_path: str, in_channels: int, num_classes: int
|
|
) -> torch.nn.Module:
|
|
"""
|
|
Build and return a PyTorch model loaded from a checkpoint.
|
|
"""
|
|
model = RFClassifier(
|
|
model=mobilenetv3(
|
|
model_size="mobilenetv3_small_050",
|
|
num_classes=num_classes,
|
|
in_chans=in_channels,
|
|
)
|
|
)
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
checkpoint = torch.load(ckpt_path, weights_only=True, map_location=device)
|
|
model.load_state_dict(checkpoint["state_dict"])
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def evaluate_checkpoint(ckpt_path: str):
|
|
"""
|
|
Loads the model from checkpoint and evaluates it on a validation set.
|
|
Prints classification metrics and plots a confusion matrix.
|
|
"""
|
|
|
|
# Load validation data
|
|
X_val, y_true, class_names = load_validation_data()
|
|
num_classes = len(class_names)
|
|
in_channels = X_val.shape[1]
|
|
|
|
# Load model
|
|
model = build_model_from_ckpt(
|
|
ckpt_path, in_channels=in_channels, num_classes=num_classes
|
|
)
|
|
|
|
# Inference
|
|
y_pred = []
|
|
with torch.no_grad():
|
|
for x in X_val:
|
|
x_tensor = torch.tensor(x[np.newaxis, ...], dtype=torch.float32)
|
|
logits = model(x_tensor)
|
|
pred = torch.argmax(logits, dim=1).item()
|
|
y_pred.append(pred)
|
|
|
|
# Print classification report
|
|
print("\nClassification Report:")
|
|
print(classification_report(y_true, y_pred, target_names=class_names, zero_division=0))
|
|
|
|
# Plot confusion matrix
|
|
plot_confusion_matrix(
|
|
y_true=np.array(y_true),
|
|
y_pred=np.array(y_pred),
|
|
classes=class_names,
|
|
normalize=True,
|
|
title="Normalized Confusion Matrix",
|
|
)
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
settings = get_app_settings()
|
|
ckpt_path = os.path.join("checkpoint_files", "inference_recognition_model.ckpt")
|
|
evaluate_checkpoint(ckpt_path)
|