You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
This repository was archived by the owner on May 28, 2026. It is now read-only.
Prototypical network wrapper for few-shot (stgcn_proto)
The unified build_model(cfg) factory in __init__.py dispatches on cfg.approach to the correct model builder. All build_*_model() functions take a Config object and return an nn.Module.
Training (src/training/)
Module
Key Functions
Imports From
config.py
Config (dataclass with embedding_dim, gcn_channels, etc.), load_config(), save_config()
(none — leaf dependency)
train.py
main() — CLI dispatcher
config, train_ce, train_prototypical
train_ce.py
train_one_epoch(), validate(), main() — cross-entropy training with label smoothing, mixup, OneCycleLR/cosine scheduler, and optional auxiliary branch losses
config, augment, dataset, models
train_prototypical.py
train_prototypical() — episodic prototypical training loop
FrameBuffer, MotionDetector, LivePredictor, ASLDisplay, run_demo() — FPS-independent motion detection via time-based velocity (displacement/dt), pre-trigger buffering (retains sign onset frames on IDLE→SIGNING), pose quality gate (skips inference if <30% valid shoulder frames), and "No body detected" display warning
config, preprocess, models
export_onnx.py
export_to_onnx(), verify_onnx(), benchmark_onnx()
config, models
Data Flow Diagrams
Training Data Flow
data/raw/*.mp4
│
v
┌─────────────────────────────────┐
│ preprocess.py │
│ extract_keypoints_mediapipe() │
│ normalize_keypoints() │ ──> data/processed/*.npy (T, 543, 3)
│ 1. shoulder-center + scale │
│ 2. face-center relative │
│ 3. depth normalization │
│ 4. hand-relative to wrist │
│ create_splits() │ ──> data/splits/WLASL{N}/train.csv
└─────────────────────────────────┘
data/processed/*.npy + data/splits/WLASL{N}/train.csv
│
v
┌─────────────────────────────────┐
│ dataset.py │
│ WLASLKeypointDataset │
│ __getitem__(): │
│ load .npy │
│ slice to 75 keypoints │ (drop face landmarks)
│ pad/crop to T frames │ (reflection padding)
│ compute velocity (motion) │ ──> (T, 75*6) when use_motion=True
│ apply augmentations │ (incl. KeypointYawRotation for 3D viewpoint simulation)
│
│ flatten to (T, input_dim) │
└─────────────────────────────────┘
│
v
┌─────────────────────────────────┐
│ train.py │
│ train_one_epoch(): │
│ mixup (if enabled) │
│ forward pass through model │
│ loss + backprop │
│ validate(): │
│ forward pass (no augment) │
│ compute top-1 / top-5 acc │ ──> checkpoints/best_model.pt
│ early stopping check │ ──> logs/ (TensorBoard)
└─────────────────────────────────┘
Inference Data Flow
Single Video (predict.py):
video.mp4 ──> MediaPipe ──> normalize (shoulder+hand-relative) ──> velocity ──> model ──> top-5 predictions
│ ^
│ OR │
keypoints.npy ──────────────> velocity ──────────────────┘
Live Demo (live_demo.py):
Webcam ──> MediaPipe ──> MotionDetector ──> FrameBuffer(max_sign_duration * camera_fps + 10)
│ │ │
│ FPS-independent: │
│ velocity = displacement/dt │
│ (dt via time.monotonic()) │
│ │ │
│ │ IDLE→SIGNING: trim buffer to pre_sign_frames (keep sign onset)
│ │ state: IDLE/SIGNING/COMPLETED
│ │ │
│ │ inference only on COMPLETED:
│ │ ├─ settle_time elapsed ──> COMPLETED
│ │ └─ max_sign_duration ──> COMPLETED
│ │ │
│ v v
│ pose quality gate: skip if <30% of buffer
│ frames have valid shoulder landmarks
│ │ (pass)
│ v
│ normalize ──> TemporalCrop(T) ──> model ──> prediction
│ (full sign frames: TemporalCrop uniformly
│ samples to T, matching training pipeline)
│ │
v v
Display <───── overlay predicted gloss + confidence + motion state
(high-conf: full cooldown, low-conf: 30% cooldown)
Shows "No body detected" warning when pose_landmarks is None
ONNX Export (export_onnx.py):
checkpoint ──> load model ──> torch.onnx.export() ──> model.onnx
│
verify (optional) ──> ONNX Runtime forward pass
benchmark (optional) ──> avg latency over 100 runs
Model Architecture Flow (ST-GCN)
Input: (batch, T, input_dim) input_dim = 75*3 or 75*6 (with motion)
│
v
┌──────────────────────────┐
│ Reshape to graph │ (B, C, T, V) where V=num_keypoints
└─────────┬────────────────┘
│
┌─────────┼─────────┐
v v v
┌──────┐ ┌──────┐ ┌──────┐
│ Body │ │ Left │ │Right │ Separate graph convolution branches
│ GCN │ │ Hand │ │ Hand │ with dilated TCN, DropPath, joint importance
│(33kp)│ │(21kp)│ │(21kp)│
└──┬───┘ └──┬───┘ └──┬───┘
└─────────┼─────────┘
v
┌─────────────────────┐
│ Avg Pool or │ Pool over time+joints per branch
│ AttentionPool │ (optional attention-weighted temporal pooling)
└─────────┬───────────┘
v
┌─────────────────────┐
│ CrossBranchAttn? │ Optional cross-branch attention fusion
└─────────┬───────────┘
v
┌─────────────────────┐
│ Concat + Project │ Fuse branch outputs
└─────────┬───────────┘
v
┌─────────────────────┐
│ L2 Normalize? │ Only when normalize_embeddings=True (proto)
└─────────┬───────────┘
v
┌─────────────────────────────┐
│ Classification head (CE) │ Linear→LayerNorm→ReLU→Dropout→Linear
│ OR Prototypical distance │ Distance to class prototypes
└─────────┬───────────────────┘
v
Output: (batch, num_classes) logits or distances
Configuration Flow
configs/stgcn_ce.yaml (default) ──> load_config() ──> Config dataclass
configs/stgcn_proto.yaml │
┌───────────┼───────────┐
v v v
train.py evaluate.py predict.py
live_demo.py
export_onnx.py
Config.__post_init__() auto-derives (scales with variant size):
wlasl_variant: 100 ──> num_classes: 100, d_model: 128, nhead: 4, num_layers: 2, dropout: 0.1
wlasl_variant: 300 ──> num_classes: 300, d_model: 192, nhead: 6, num_layers: 4, dropout: 0.3
wlasl_variant: 1000 ──> num_classes: 1000, d_model: 256, nhead: 8, num_layers: 5, dropout: 0.4
wlasl_variant: 2000 ──> num_classes: 2000, d_model: 384, nhead: 8, num_layers: 6, dropout: 0.5
Note: d_model, nhead, num_layers, dropout are auto-scaled per variant.
YAML-configurable model fields (read by load_config into Config dataclass):
embedding_dim: 128 # Final embedding dimension for ST-GCN encoder
gcn_channels: [64, 128, 128] # Channel widths per ST-GCN block
Tests (tests/)
Each test file maps to one or more source modules:
Test File
Tests For
test_augment.py
src/data/augment.py — all transform classes and pipeline presets