added doc

2026-10-06 16:37:12 +02:00
parent 3c8123ae95
commit 787e2f47ec
5 changed files with 1301 additions and 343 deletions
@@ -1,325 +0,0 @@
 # YOLO face detector
 A small single-shot face detector in the style of YOLO v1. It splits the image
 into a 6 × 6 grid and predicts at most one face per grid cell. The website uses
 it for live webcam face detection over a websocket.
 | | |
 |---|---|
 | Model code | [`api/src/models/yolo.py`](../src/models/yolo.py) (`Yolo_model`, `SkipBlock`, `convert_prediction`) |
 | Training code | [`api/src/yolo/`](../src/yolo/) (`config.py`, `dataset.py`, `loss.py`, `tasks.py`) |
 | Deployed weights | `api/build_models/face_detection_yolo` (PyTorch `state_dict`) |
 | Production use | [`api/src/production/inference.py`](../src/production/inference.py) `detect_faces`, websocket `/predict_face` |
 | Input | RGB image resized to 64 × 64, values in [0, 1] |
 | Output | 6 × 6 grid × 6 values per cell (x, y, w, h, confidence, class) |
 | Classes | 1 (face) |
 | Parameters | 657,254 (with `c_hidden=16`) |
 | Speed | about 18 ms per image on CPU, batch size 1 (measured on the dev machine) |
 ## Architecture
 The network is a fully convolutional stack of residual "skip blocks". It never
 downsamples: every layer works at the full 64 × 64 resolution. At the end, one
 `AdaptiveAvgPool2d` reduces the feature map to the 6 × 6 grid, and a 1 × 1
 convolution turns each grid cell's features into a prediction.
 ```mermaid
 flowchart TD
    IN["Input image<br/>3 × 64 × 64"]
    subgraph STEM["Stem"]
        S1["Conv 3×3, 3 → 16<br/>GroupNorm(2) · LeakyReLU"]
    end
    subgraph ST1["Stage 1 · 16 channels"]
        B1["SkipBlock 16 → 16"] --> B2["SkipBlock 16 → 16"] --> B3["SkipBlock 16 → 16"] --> B4["SkipBlock 16 → 16"]
    end
    subgraph ST2["Stage 2 · 32 channels"]
        B5["SkipBlock 16 → 32"] --> B6["SkipBlock 32 → 32"] --> B7["SkipBlock 32 → 32"] --> B8["SkipBlock 32 → 32"]
    end
    subgraph ST3["Stage 3 · 64 channels"]
        B9["SkipBlock 32 → 64"] --> B10["SkipBlock 64 → 64"] --> B11["SkipBlock 64 → 64"] --> B12["SkipBlock 64 → 64"]
    end
    subgraph HEAD["Head"]
        H1["Conv 3×3, 64 → 128<br/>GroupNorm(8) · LeakyReLU · Dropout 0.3"]
        H2["AdaptiveAvgPool2d → 6 × 6<br/>128 × 6 × 6"]
        H3["Conv 1×1, 128 → 6<br/>6 × 6 × 6"]
        H1 --> H2 --> H3
    end
    ACT["permute to [B, 6, 6, 6]<br/>activation per channel"]
    OUT["Per cell: x · y · w · h · conf · class"]
    IN --> STEM --> ST1 --> ST2 --> ST3 --> HEAD --> ACT --> OUT
 ```
 Every layer up to and including the head's 3 × 3 conv keeps the spatial size at
 64 × 64. Only the channel count changes.
 | Index in `model` | Layer | Output shape (batch 1) | Parameters |
 |---|---|---|---|
 | 0–2 | Conv 3×3 + GroupNorm + LeakyReLU | 16 × 64 × 64 | 480 |
 | 3–6 | 4 × SkipBlock 16 → 16 | 16 × 64 × 64 | 4 × 7,360 |
 | 7–10 | SkipBlock 16 → 32, 3 × SkipBlock 32 → 32 | 32 × 64 × 64 | 23,936 + 3 × 29,056 |
 | 11–14 | SkipBlock 32 → 64, 3 × SkipBlock 64 → 64 | 64 × 64 × 64 | 94,976 + 3 × 115,456 |
 | 15–18 | Conv 3×3 + GroupNorm + LeakyReLU + Dropout | 128 × 64 × 64 | 74,112 |
 | 19 | AdaptiveAvgPool2d(6, 6) | 128 × 6 × 6 | 0 |
 | 20 | Conv 1×1 | 6 × 6 × 6 | 774 |
 The channel counts scale with `c_hidden` (default 16): c, 2c, 4c, then 8c in the
 head. The output has `boxes * 5 + labels` channels. With 1 box and 1 label
 that is 6.
 ### SkipBlock
 Each block is a residual unit. Three 3 × 3 convolutions form the main path and
 a 1 × 1 convolution forms the shortcut. The 1 × 1 conv lets the shortcut change
 the channel count.
 ```mermaid
 flowchart LR
    X["x<br/>C_in × 64 × 64"]
    subgraph MAIN["main path"]
        C1["Conv 3×3 → C_out<br/>GroupNorm · LeakyReLU"]
        C2["Conv 3×3<br/>GroupNorm · LeakyReLU"]
        C3["Conv 3×3<br/>GroupNorm · LeakyReLU"]
        C1 --> C2 --> C3
    end
    subgraph SKIP["shortcut"]
        K["Conv 1×1 → C_out<br/>GroupNorm · LeakyReLU"]
    end
    ADD(("+"))
    POST["LeakyReLU<br/>Dropout p=0.3 (always on)"]
    Y["y<br/>C_out × 64 × 64"]
    X --> C1
    X --> K
    C3 --> ADD
    K --> ADD
    ADD --> POST --> Y
 ```
 All GroupNorm layers use `C / 8` groups, so each group holds 8 channels.
 ### Design notes
 - **No downsampling.** All twelve blocks run at 64 × 64. This makes the network
  cheap in parameters but comparatively expensive in compute and memory for its
  size.
 - **Receptive field.** The network has 38 stacked 3 × 3 convolutions (1 stem,
  36 in the blocks, 1 in the head). That gives a receptive field of
  1 + 38 × 2 = 77 px, so every output cell can see the whole 64 px image.
 - **Uneven pooling bins.** 64 is not divisible by 6, so `AdaptiveAvgPool2d`
  averages bins of 10–11 px, and neighbouring bins overlap slightly.
 - **Layer names are part of the file format.** `build_models/face_detection_yolo`
  is a plain `state_dict` keyed by names like `model.3.conv.0.weight`. Adding,
  removing or reordering layers in `Yolo_model` breaks loading of the deployed
  weights.
 ## Output encoding
 `forward` returns a tensor of shape `[batch, 6, 6, 6]`. For each grid cell
 `[i, j]` the last dimension holds:
 | Channel | Meaning | Range | Activation |
 |---|---|---|---|
 | 0 | `x`: box centre inside the cell, as a fraction of the cell width | 0–1 | sigmoid |
 | 1 | `y`: box centre inside the cell, as a fraction of the cell height | 0–1 | sigmoid |
 | 2 | `w`: box width as a fraction of the image width | ≥ 0 | `exp` or sigmoid |
 | 3 | `h`: box height as a fraction of the image height | ≥ 0 | `exp` or sigmoid |
 | 4 | confidence that a face's centre lies in this cell | 0–1 | sigmoid |
 | 5 | class score (face) | 0–1 | sigmoid |
 `size_activation` decides how `w` and `h` are produced:
 - `"exp"`: `w = exp(raw)`, unbounded. The **deployed** weights were trained
  this way, so production sets `FACE_SIZE_ACTIVATION = "exp"`.
 - `"sigmoid"`: `w = sigmoid(raw)`, bounded to 0–1. This is the **training
  default** in `Config.size_activation`.
 A model must always run with the activation it was trained with.
 **Grid index convention.** The first grid index `i` is the horizontal cell
 (column) and the second `j` is the vertical cell (row). The dataset builds
 targets this way, and `convert_prediction` reads them back the same way.
 The convolutional feature map is laid out `[row, column]`, so the network
 effectively learns a spatial transpose. Its receptive field covers the whole
 image, so it can do that. A check on the deployed model confirms it: on 69
 validation images with one face in an off-diagonal cell, the most confident
 cell matched the `[column, row]` target 49 times and the transposed cell once.
 If you change the dataset or the decoding, keep both on the same convention.
 ## From prediction to boxes
 `convert_prediction(label, image, threshold)` turns one `[6, 6, 6]` output into
 pixel boxes. It also works on a target tensor, which is how `view-data` draws
 the ground truth.
 ```mermaid
 flowchart LR
    P["Prediction<br/>6 × 6 × 6"] --> F{"conf ≥ threshold?"}
    F -- no --> E["empty cell<br/>(only used for drawing)"]
    F -- yes --> D["decode cell i, j<br/>to centre + size in px"]
    D --> C["corners<br/>xmin, ymin, xmax, ymax"]
    C --> N["NMS<br/>IoU threshold 0.1"]
    N --> S["scale to the original<br/>image size (production)"]
 ```
 For a cell `[i, j]` on a 64 px image with a 6 × 6 grid (cell size
 `c = 64 / 6 ≈ 10.67 px`):
 ```
 centre_x = (i + x) · c          width  = w · 64
 centre_y = (j + y) · c          height = h · 64
 xmin = centre_x − width / 2     xmax = xmin + width     (width and height are at least 1 px)
 ymin = centre_y − height / 2    ymax = ymin + height
 ```
 Non-maximum suppression (`nms` in
 [`api/src/common/boxes.py`](../src/common/boxes.py)) sorts the boxes by
 confidence. It keeps the best box and drops every other box that overlaps it
 with IoU ≥ 0.1. The low threshold suits faces, which rarely overlap.
 In production, `detect_faces` multiplies the boxes by `W / 64` and `H / 64` to
 map them back to the size of the uploaded frame.
 ## Training
 ### Data
 - **Format.** COCO: images plus `_annotations.coco.json` in each folder.
  Defaults: `api/data/faces_scenery/train` and `api/data/faces_scenery/test`.
 - **Resize.** Images are resized to 64 × 64.
 - **Augmentation** (albumentations):
  - horizontal flip, p = 0.5
  - brightness/contrast, p = 0.2
  - affine: translate ±40 %, scale 0.5–1.5, edge pixels replicated, p = 0.5
  - boxes that keep less than 50 % of their area are dropped
 - **Augmented validation.** `augment_val=True` by default, so the validation
  set is augmented too. The original code did this.
 - **Targets.** `YoloDataset` builds a `[6, 6, 5 + classes]` target per image.
  - Each cell takes the largest face whose centre lies in it. Annotations are
    sorted by area, largest first, and each face is assigned only once.
  - `x` and `y` are the centre's offset inside the cell, divided by the cell
    size.
  - `w` and `h` are the box size divided by the image size.
  - Confidence is 1, and the class is one-hot.
  - All other cells are zero.
 ### Loss
 `YoloLoss` ([`api/src/yolo/loss.py`](../src/yolo/loss.py)) is the sum of
 four unweighted terms:
 | Term | Cells | Function |
 |---|---|---|
 | box | cells with a face | MSE on (x, y, w, h) |
 | object | cells with a face | MSE on confidence vs 1 |
 | no-object | empty cells | MSE on confidence vs 0 |
 | class | cells with a face | binary cross-entropy (with logits) |
 Each term is averaged over its own set of cells. So the roughly 35 empty cells
 together weigh the same as the one or two face cells, without the λ factors the
 original YOLO paper uses.
 ### Running
 The commands run from `api/`. Every `Config` field can be overridden on the
 command line.
 ```bash
 python -m src.main yolo.train                    # train (500 epochs, batch 64, Adam lr 1e-3, wd 1e-4)
 python -m src.main yolo.train --epochs 50 --size-activation exp
 python -m src.main yolo.view-data                # show validation images with their target grid
 python -m src.main yolo.sample                   # show the trained model's predictions
 python -m src.main yolo.webcam --camera 0        # live detection on a local webcam, q quits
 python -m src.main yolo.benchmark-workers        # find the fastest DataLoader num_workers
 python -m src.main yolo.train --help             # list every option
 ```
 - **Checkpoints.** Training keeps the checkpoint with the lowest validation
  loss at `saved_models/face_detection_yolo/face_detection_yolo`.
 - **Drawn grids.** `sample`, `view-data` and `webcam` draw red boxes for
  detections, blue for cells with a face, and green for empty cells.
 - **Thresholds.** These tasks draw a box from confidence 0.95 (`--threshold`).
  Production uses 0.9.
 ### Deploying a new model
 1. Train it; the result is in `api/saved_models/face_detection_yolo/`.
 2. Copy that file to `api/build_models/face_detection_yolo`.
 3. Make the settings in
   [`api/src/production/settings.py`](../src/production/settings.py) match
   how it was trained:
   - `FACE_SIZE_ACTIVATION`: `"sigmoid"` if you kept the training default.
   - `FACE_C_HIDDEN`, `FACE_GRID` and `FACE_IMAGE_SIZE`.
 4. Restart the API. A size mismatch fails at startup in `load_checkpoint`.
   An activation mismatch does not fail; it silently produces wrong box sizes.
 ## Use in production
 ```mermaid
 sequenceDiagram
    participant B as Browser (Object_Detection.jsx)
    participant W as /predict_face websocket
    participant I as detect_faces
    participant M as Yolo_model
    B->>W: webcam frame as JPEG bytes
    Note over W: checks origin, connection limits,<br/>frame size and FACE_MAX_FPS.<br/>Frames that arrive too fast or<br/>while busy are dropped
    W->>I: frame bytes
    I->>I: decode, resize to 64 × 64
    I->>M: tensor [1, 3, 64, 64]
    M-->>I: [1, 6, 6, 6]
    I->>I: convert_prediction (threshold 0.9) + NMS,<br/>scale to frame size
    I-->>W: boxes
    W-->>B: {"bboxes": [[xmin, ymin, xmax, ymax], ...]}
    B->>B: draw boxes on the overlay canvas
 ```
 Relevant settings in `api/src/production/settings.py`:
 | Setting | Value | Meaning |
 |---|---|---|
 | `FACE_MODEL_FILE` | `face_detection_yolo` | file in `api/build_models/` |
 | `FACE_IMAGE_SIZE` | 64 | model input size |
 | `FACE_GRID` | 6 | grid size |
 | `FACE_C_HIDDEN` | 16 | base channel count |
 | `FACE_SIZE_ACTIVATION` | `"exp"` | must match the weights |
 | `FACE_THRESHOLD` | 0.9 | minimum confidence for a box |
 | `FACE_MAX_FPS` | 5.5 (env) | frames per second processed per websocket |
 ## Known quirks and limitations
 - **Dropout is always on.** `SkipBlock.forward` calls `F.dropout(..., p=0.3)`,
  which defaults to `training=True`, so it stays active in `eval()` mode.
  - Two runs on the same image give slightly different outputs, and boxes can
    flicker between video frames.
  - The deployed weights were trained this way, so it is left as is.
  - The head's `nn.Dropout` behaves normally and is off in eval mode.
 - **Activation default differs from the deployed model.** Training defaults to
  `"sigmoid"`, while the deployed weights need `"exp"`. See
  [Deploying a new model](#deploying-a-new-model).
 - **Class loss applies a sigmoid twice.** The model already outputs sigmoid
  values, and `binary_cross_entropy_with_logits` applies another sigmoid. The
  class output is unused with a single class (`convert_prediction` ignores it),
  so this does not affect detections.
 - **One face per cell, at most 36 faces.** When two face centres fall into the
  same cell, only the larger one is a training target, and the model can
  predict only one box there.
 - **No anchors, one box per cell.** `boxes > 1` is accepted by the constructor,
  but the dataset, loss and `convert_prediction` only handle the first box.
 - **Small input.** At 64 × 64, faces that are a few pixels tall after resizing
  are hard to detect. Very wide or tall frames are squashed, because the
  resize does not keep the aspect ratio.
 - **Noisy validation loss.** The validation set is augmented by default, so the
  loss that picks the best checkpoint changes from epoch to epoch even with
  fixed weights.
@@ -60,17 +60,26 @@ class Yolo_model(nn.Module):
            SkipBlock(c_in=c_hidden, c_out=c_hidden),
            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden * 2),
+            nn.Conv2d(c_hidden, c_hidden * 2, kernel_size=3, padding=1, stride=2),
            nn.LeakyReLU(inplace=True),
            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
-            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 4),
+            nn.Conv2d(c_hidden * 2, c_hidden * 4, kernel_size=3, padding=1, stride=2),
            nn.LeakyReLU(inplace=True),
            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
-            nn.Conv2d(c_hidden * 4, c_hidden * 8, kernel_size=3, padding=1),
+            nn.Conv2d(c_hidden * 4, c_hidden * 8, kernel_size=3, padding=1, stride=2),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(c_hidden * 8, c_hidden * 8, kernel_size=3, padding=1),
            nn.BatchNorm2d(c_hidden * 8),
            nn.LeakyReLU(inplace=True),
@@ -13,41 +13,39 @@ class YoloLoss(nn.Module):
        target_boxes = targets[..., :4]
        target_conf = targets[..., 4]
        target_classes = targets[..., 5:]
        N = predictions.shape[0]
        obj_mask = targets[..., 4] == 1
        noobj_mask = targets[..., 4] == 0
-        if obj_mask.any():
+        box_loss = F.mse_loss(
-            box_loss = F.mse_loss(
+            pred_boxes[obj_mask],
-                pred_boxes[obj_mask],
+            target_boxes[obj_mask],
-                target_boxes[obj_mask],
+            reduction="sum"
-                reduction="mean"
+        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
            )
        else:
            box_loss = torch.tensor(0.0, device=predictions.device)
        obj_loss = F.mse_loss(
            pred_conf[obj_mask],
            target_conf[obj_mask],
-            reduction="mean"
+            reduction="sum"
        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
        noobj_loss = F.mse_loss(
            pred_conf[noobj_mask],
            target_conf[noobj_mask],
-            reduction="mean"
+            reduction="sum"
        ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device)
        # The model already applies a sigmoid, so these are probabilities.
        class_loss = F.binary_cross_entropy(
            pred_classes[obj_mask],
-            target_classes[obj_mask]
+            target_classes[obj_mask],
            reduction="sum"
        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
-        total_loss = box_loss + obj_loss + noobj_loss + class_loss
+        total_loss = (box_loss * 5 + obj_loss + noobj_loss * 0.5 + class_loss) / N
        return total_loss
@@ -3,8 +3,8 @@
 for ((c = 1; c <= 32; c *= 2)); do
    python -m src.main yolo.train --train-dir ./data/faces_scenery/train/ --val-dir ./data/faces_scenery/test/ \
-        --no-augment-val --c-hidden "$c" --batch-size 64 --epochs 200 --run-name "bn-c$c" \
+        --no-augment-val --c-hidden "$c" --batch-size 64 --epochs 200 --run-name "bn-weighted-loss-c$c" \
-        --no-save --train-workers 2 --image-size 64 --grid 6
+        --no-save --train-workers 4 --image-size 64 --grid 6
 done