File size: 1,708 Bytes
215f97f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """
Training utilities — class weighting and label encoding helpers.
Class weighting is critical: with ~20-25% positives in the training set,
the model will bias toward predicting ``"safe"`` without it.
"""
import numpy as np
from sklearn.utils.class_weight import compute_class_weight
from app.ml.cnn.architecture import LABEL_NAMES
from app.core.logging import get_logger
logger = get_logger(__name__)
def get_class_weights(labels_onehot: np.ndarray) -> dict[int, float]:
"""Compute balanced class weights from one-hot encoded labels.
Uses ``sklearn.utils.class_weight.compute_class_weight`` with
``class_weight="balanced"`` to inversely weight classes by frequency.
Args:
labels_onehot: One-hot encoded labels, shape ``(N, 3)``.
Returns:
Dict mapping class index → weight, suitable for
``model.fit(..., class_weight={"label": weights})``.
"""
# Convert one-hot back to integer labels
y_int = np.argmax(labels_onehot, axis=1)
classes = np.arange(len(LABEL_NAMES))
# Calculate manually to avoid sklearn's ValueError if a class is entirely missing (e.g. during bootstrap)
total_samples = len(y_int)
num_classes = len(classes)
weight_dict = {}
for cls in classes:
cls_count = np.sum(y_int == cls)
if cls_count > 0:
weight = total_samples / (num_classes * cls_count)
else:
weight = 1.0 # default weight for missing classes
weight_dict[int(cls)] = float(weight)
logger.info(
"Class weights: %s",
{LABEL_NAMES[k]: f"{v:.3f}" for k, v in weight_dict.items()},
)
return weight_dict
|