added new metrics

api/src/yolo/loss.py

@@ -41,7 +41,8 @@ class YoloLoss(nn.Module):
         ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
 
-        class_loss = F.binary_cross_entropy_with_logits(
+        # The model already applies a sigmoid, so these are probabilities.
+        class_loss = F.binary_cross_entropy(
             pred_classes[obj_mask],
             target_classes[obj_mask]
         ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)

api/src/yolo/metrics.py

@@ -0,0 +1,51 @@
+"""Detection metrics over the whole validation set, logged next to the loss.
+Per grid cell, like the loss: a cell is a detection when its confidence
+passes `conf_threshold`, and a correct one when the cell holds a face and
+the box overlaps the true box by at least `iou_threshold`."""
+
+import torch
+
+
+def _cell_boxes(t):
+    """[N, S, S, >=4] cell outputs (x, y relative to the cell, w, h relative to
+    the image) -> [N, S, S, 4] boxes as (xmin, ymin, xmax, ymax) in image units."""
+    S = t.shape[1]
+    i = torch.arange(S, device=t.device).view(S, 1)
+    j = torch.arange(S, device=t.device).view(1, S)
+    cx = (i + t[..., 0]) / S
+    cy = (j + t[..., 1]) / S
+    w, h = t[..., 2], t[..., 3]
+    return torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], dim=-1)
+
+
+def _iou(a, b):
+    lt = torch.max(a[..., :2], b[..., :2])
+    rb = torch.min(a[..., 2:], b[..., 2:])
+    inter = (rb - lt).clamp(min=0).prod(-1)
+    area_a = (a[..., 2:] - a[..., :2]).clamp(min=0).prod(-1)
+    area_b = (b[..., 2:] - b[..., :2]).clamp(min=0).prod(-1)
+    return inter / (area_a + area_b - inter).clamp(min=1e-9)
+
+
+def _average_precision(scores, labels):
+    """How well the confidence ranks cells with a face above empty cells (1 = perfect)."""
+    order = scores.argsort(descending=True)
+    hits = labels[order].float()
+    precision = hits.cumsum(0) / torch.arange(1, len(hits) + 1, device=hits.device)
+    return ((precision * hits).sum() / hits.sum().clamp(min=1)).item()
+
+
+def detection_metrics(preds, targets, conf_threshold=0.5, iou_threshold=0.5):
+    obj = targets[..., 4] == 1
+    conf = preds[..., 4]
+    ious = _iou(_cell_boxes(preds), _cell_boxes(targets))
+
+    detected = conf > conf_threshold
+    correct = (detected & obj & (ious > iou_threshold)).sum().item()
+    n_obj = obj.sum().item()
+    return {
+        "conf_ap": _average_precision(conf.flatten(), obj.flatten()),
+        "box_iou": ious[obj].mean().item() if n_obj else 0.0,
+        "precision": correct / max(detected.sum().item(), 1),
+        "recall": correct / max(n_obj, 1),
+    }

api/src/yolo/tasks.py

@@ -16,6 +16,7 @@ from src.models.yolo import Yolo_model, convert_prediction
 from src.yolo.config import Config
 from src.yolo.dataset import YoloDataset
 from src.yolo.loss import YoloLoss
+from src.yolo.metrics import detection_metrics
 
 
 def build_dataset(cfg: Config, folder: str, augment: bool):
@@ -62,7 +63,7 @@ def train(cfg: Config):
 
     fit(model, YoloLoss(), optimizer, train_loader, val_loader, device,
         epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
-        monitor="loss", config=cfg)
+        monitor="loss", val_metrics=detection_metrics, config=cfg)
 
 
 def view_data(cfg: Config):