A Diffusion Model (DDPM + DDIM) that generates realistic defensive player movements conditioned on an offensive play sequence, using ResBlock1D + Self-Attention + Cross-Attention.
git clone https://github.com/TQG1997/basketball.git
cd basketball
pip install -r requirements.txt
# Download dataset (~730MB)
pip install gdown
gdown --folder https://drive.google.com/drive/folders/1uNPw7LOA3xENclQRtSlUftiR7tlVNOts -O data/
# Train (auto-tunes batch/filters for your GPU)
python src/train.py --data_path=data --output=output --max_epochs=2000 --auto --yes
One-click training on free GPU. Use Runtime → Change runtime type → T4 GPU.
Native macOS/Linux desktop app for sketching offensive plays and visualizing AI-generated defensive responses:
pip install -r requirements-qt.txt
cd ui && python Main.pyRequires a trained checkpoint at ui/Data/checkpoints/model_epoch500.pt and normalization data in ui/Data/Model_data/.
┌──────────────────────────────┬──────────────────────────────┐
│ 🖼️ Drawing Board (left) │ 🎬 Animation (right) │
│ │ │
│ Double-click → ball (+auto draw)│ Radio: Sketch Animation │
│ Drag PG/SG/SF/PF/C to position │ Play Simulation │
│ Click-drag → draw ball path │ │
│ Undo / Generate → AI defense │ ▶ Play/Pause ─●─ slider │
└──────────────────────────────┴──────────────────────────────┘
Step-by-step:
| Step | Action | What Happens |
|---|---|---|
| ① | Double-click court (left) | Ball appears. Nearest player auto-snaps. Draw mode ON. |
| ② | Drag players (PG/SG/SF/PF/C) | Reposition offensive formation on the court. |
| ③ | Click & drag on court | Draw ball trajectory — gold dotted arrow shows pass/shot. |
| ④ | Repeat ③ for each pass/dribble | Build multi-segment play: pass→SG→dribble→pass→C→shoot. |
| ⑤ | Undo (if mistake) | Removes last segment without resetting everything. |
| ⑥ | Click → Generate | AI generates defensive play. Status: Saving → Generating → Done! |
| ⑦ | Sketch Animation | Replay your offensive play sketch. |
| ⑧ | Play Simulation | Watch AI-generated defensive response. ▶/⏸ + slider. |
| ⑨ | Clear | Reset all to start fresh. |
Tips:
- Ball near basket = shot attempt (auto-detected). Ball near player = pass.
- Double-click a moving player to add wiggle animation to their path.
- PG=Point Guard, SG=Shooting Guard, SF=Small Forward, PF=Power Forward, C=Center.
Basketball/
├── src/ # PyTorch diffusion model + training
│ ├── train.py # Training entry (AMP, EMA, validation)
│ ├── diffusion.py # DDPM/DDIM + DenoiserNet
│ ├── ops.py # Conv1D_SN, ResBlock1D, Self/CrossAttention
│ ├── game_visualizer.py # Play animation (matplotlib)
│ └── utils.py # DataFactory re-export
├── ui/ # Desktop app + Web UI + inference
│ ├── Main.py # PyQt5 desktop app (main entry)
│ ├── Drawingboard.py # Interactive sketch canvas
│ ├── Court.py # Animation playback
│ ├── Players.py / Ball.py # Draggable player & ball items
│ ├── SavePos.py / Bezier.py # Trajectory smoothing
│ ├── app.py # Gradio web interface (alternative)
│ ├── inference.py # PyTorch inference pipeline
│ └── draw_feat.py # Ball-possession feature extraction
├── shared/ # DataFactory (numpy, framework-agnostic)
├── config/ # Training YAML configuration
├── notebooks/ # Colab training notebook
├── data/ # Dataset (.npy, excluded from git)
└── requirements.txt
| Parameter | Default | Description |
|---|---|---|
--data_path |
data |
Dataset directory |
--output |
output |
Output dir (checkpoints, samples) |
--max_epochs |
null |
Stop after N epochs (null = forever) |
--resume |
false |
Resume from latest checkpoint in output dir |
--batch_size |
64 |
Batch size (overridden by --auto) |
--lr |
1e-4 |
AdamW learning rate |
--n_filters |
256 |
Conv filters |
--n_resblock |
4 |
Residual blocks per network |
--num_heads |
4 |
Attention heads |
--T |
1000 |
Diffusion timesteps |
--ddim_steps |
50 |
DDIM sampling steps |
--checkpoint_step |
100 |
Epochs between checkpoints |
--vis_freq |
10 |
Epochs between visualizations |
--auto |
false |
Auto-tune batch/filters for GPU VRAM (80% target) |
--yes |
false |
Skip output-dir confirmation prompt |
Use --auto for automatic VRAM-aware configuration:
| GPU | VRAM | --auto sets |
|---|---|---|
| T4 (Colab free) | 15 GB | batch≈350, filters=1536, resblocks=10 |
| A10G | 24 GB | batch=512, filters=2048, resblocks=14 |
| A100 | 40 GB | batch=512, filters=2048, resblocks=14 |
Override any auto setting: --auto --batch_size=128. See config/config.yaml for all settings.
Offensive Play (conditioning) Random Noise
│ │
▼ ▼
┌──────────┐ ┌──────────────┐
│ Cond Proj │ │ Input Proj │
└──────────┘ └──────────────┘
│ │
└─────────────┬─────────────────────┘
▼
┌───────────────────┐
│ ResBlock1D × N │ (LayerNorm + SiLU + Conv1D)
│ Self-Attention │ (temporal interactions)
│ Cross-Attention │ (attend to conditioning)
└───────────────────┘
│
▼
┌──────────┐
│ Output │ → Defence (10) + Ball Features (6)
└──────────┘
Training: DDPM forward process → predict noise → MSE loss
Inference: DDIM reverse process — 50 steps from noise to trajectory
Input (Conditioning): Output (Generated):
┌──────────────────────┐ ┌──────────────────────┐
│ Offence [T×12] │ │ Defence [T×10] │
│ ball x,y │ │ B1..B5 x,y │
│ A1..A5 x,y │ → │ │
│ │ │ Ball Features [T×6] │
│ Seq Feat [T×6] │ │ dribble/pass status │
│ ball possession │ │ │
└──────────────────────┘ └──────────────────────┘
T = 50 frames (~8 seconds of gameplay)
Conditional time-series generation: given offense, predict defense
-
Flow Matching — Replace DDPM with flow matching for faster training and sampling (10-20 steps instead of 50-1000). Reference
-
Graph Neural Network for Player Interactions — Model 10 players + ball as graph nodes with edges encoding distance and defensive matchups, replacing the current flat 16-dim vector representation. More physically grounded player movement modeling.
-
Variable-Length Sequences via Perceiver IO — Remove the fixed 50-frame limit. Handle plays of arbitrary length with attention-based pooling over time, enabling the model to generalize across different play durations.
| File | Shape | Size | Description |
|---|---|---|---|
50Real.npy |
(14032, 50, 11, 4) | 236MB | Ground truth plays (ball + 10 players) |
50Seq.npy |
(14032, 50, 12) | 64MB | Offence conditioning |
SeqCond.npy |
(14032, 50, 6) | 16MB | Ball-possession features (conditioning) |
RealCond.npy |
(14032, 50, 6) | 16MB | Ball-possession features (ground truth) |
Feature layout: entity 0 = ball (x, y, z, flag), entities 1-5 = offence A1-A5, entities 6-10 = defence B1-B5.
Download from Google Drive. Data split 9:1 train/valid by DataFactory.