# 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
3 × 64 × 64"] subgraph STEM["Stem"] S1["Conv 3×3, 3 → 16
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
GroupNorm(8) · LeakyReLU · Dropout 0.3"] H2["AdaptiveAvgPool2d → 6 × 6
128 × 6 × 6"] H3["Conv 1×1, 128 → 6
6 × 6 × 6"] H1 --> H2 --> H3 end ACT["permute to [B, 6, 6, 6]
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
C_in × 64 × 64"] subgraph MAIN["main path"] C1["Conv 3×3 → C_out
GroupNorm · LeakyReLU"] C2["Conv 3×3
GroupNorm · LeakyReLU"] C3["Conv 3×3
GroupNorm · LeakyReLU"] C1 --> C2 --> C3 end subgraph SKIP["shortcut"] K["Conv 1×1 → C_out
GroupNorm · LeakyReLU"] end ADD(("+")) POST["LeakyReLU
Dropout p=0.3 (always on)"] Y["y
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
6 × 6 × 6"] --> F{"conf ≥ threshold?"} F -- no --> E["empty cell
(only used for drawing)"] F -- yes --> D["decode cell i, j
to centre + size in px"] D --> C["corners
xmin, ymin, xmax, ymax"] C --> N["NMS
IoU threshold 0.1"] N --> S["scale to the original
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,
frame size and FACE_MAX_FPS.
Frames that arrive too fast or
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,
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.