TrackBall Lite

TrackBall Lite is an original, lightweight tennis-ball detector and trajectory tracker designed for single-class mobile inference. It is not a YOLO or TrackNet fork: the implementation uses a compact depthwise backbone, a motion-gated heatmap/box head, and a deterministic trajectory filter built for fast-moving small balls.

Goals

• Detect only moving tennis balls.
• Track ball trajectory across frames.
• Keep the model small enough for iPhone 7-class devices.
• Use an original architecture and training recipe that can be owned and modified independently.
• Keep the project runnable from zero with a simple COCO-like dataset.

Project Layout

src/trackball_lite/
  config.py        Model and training configuration dataclasses
  dataset.py       Tennis-ball dataset loader and target creation
  export.py        ONNX/Core ML export helpers
  losses.py        Heatmap, box, and motion losses
  model.py         TrackBallNet architecture
  postprocess.py   Decode heatmaps and boxes into detections
  tracker.py       Constant-velocity trajectory tracker
scripts/
  train.py         Training entry point
  infer_video.py   Video inference entry point
  export_mobile.py Mobile export entry point
tests/
  test_*.py        Fast tests for config, decoding, tracking, and export guards
docs/
  architecture.md  Design rationale and paper mapping
  dataset.md       Dataset format and labeling guide

Dataset Format

Use one JSON file per split:

{
  "images": [{"id": 1, "file_name": "frames/000001.jpg", "width": 1280, "height": 720}],
  "annotations": [{"image_id": 1, "bbox": [620, 340, 12, 12], "visible": true}]
}

Only one class is supported. The bbox is [x, y, width, height] in pixels.

Quick Start

Install runtime dependencies:

pip install -e ".[train]"

Run tests:

python -m pytest

Run a dependency-light smoke check:

$env:PYTHONPATH="src"; python scripts/smoke_test.py

Train:

python scripts/train.py --data-root data/tennis --train-json train.json --val-json val.json --epochs 80

Create clean-room synthetic smoke data and run one training epoch:

python scripts/make_synthetic_dataset.py --output data/synthetic_motion --train 18 --val 6 --sequence-length 6
$env:PYTHONPATH="src"; python scripts/train.py --data-root data/synthetic_motion --train-json train.json --val-json val.json --epochs 1 --batch-size 3 --output-dir runs/motion_smoke_train

Validate dataset structure and ownership manifest before training:

python scripts/validate_dataset.py --data-root data/synthetic_motion_big --json train.json --manifest docs/data_manifest.example.json

Export:

python scripts/export_mobile.py --checkpoint runs/best.pt --format onnx --output trackball_lite.onnx

TorchScript mobile-compatible export, useful when ONNX/Core ML dependencies are unavailable:

$env:PYTHONPATH="src"; python scripts/export_mobile.py --checkpoint runs/motion_smoke_train/best.pt --format torchscript --output runs/motion_smoke_train/trackball_lite.ptl

PyTorch Mobile Lite Interpreter export:

$env:PYTHONPATH="src"; python scripts/export_mobile.py --checkpoint runs/motion_big_train_20/best.pt --format torchscript-lite --output runs/motion_big_train_20/trackball_lite_lite.ptl
$env:PYTHONPATH="src"; python scripts/export_mobile.py --checkpoint runs/motion_big_train_20/best.pt --format torchscript-lite-optimized --output runs/motion_big_train_20/trackball_lite_lite_opt.ptl

Core ML neuralnetwork export:

$env:PYTHONPATH="src"; python scripts/export_mobile.py --checkpoint runs/motion_big_train_20/best.pt --format coreml-neuralnetwork --output runs/motion_big_train_20/trackball_lite_nn_v2.mlmodel

Evaluate detection and trajectory metrics:

$env:PYTHONPATH="src"; python scripts/evaluate.py --checkpoint runs/motion_big_train_20/best.pt --data-root data/synthetic_motion_big --json val.json --output runs/motion_big_train_20/metrics.json

Benchmark the TorchScript artifact on local CPU:

python scripts/benchmark_mobile.py --model runs/motion_big_train_20/trackball_lite.ptl --warmup 10 --iters 50 --output runs/motion_big_train_20/benchmark_cpu.json
python scripts/benchmark_mobile.py --model runs/motion_big_train_20/trackball_lite_lite.ptl --warmup 10 --iters 50 --output runs/motion_big_train_20/benchmark_cpu_lite.json

Run video inference with a TorchScript artifact and save trajectory outputs:

python scripts/make_synthetic_video.py --output runs/motion_big_train_20/synthetic_infer.mp4 --frames 36
$env:PYTHONPATH="src"; python scripts/infer_video.py --checkpoint runs/motion_big_train_20/trackball_lite.ptl --model-format torchscript --video runs/motion_big_train_20/synthetic_infer.mp4 --csv-output runs/motion_big_train_20/tracks.csv --video-output runs/motion_big_train_20/annotated_infer.mp4

Verify release artifacts:

python scripts/verify_release.py --run-dir runs/motion_big_train_20 --output runs/motion_big_train_20/release_verification.json
python scripts/verify_release.py --run-dir runs/motion_big_train_20 --distractor-run-dir runs/motion_distractors_train_20 --output runs/motion_big_train_20/release_verification.json
python scripts/verify_release.py --run-dir runs/motion_big_train_20 --strict-real-world

The distractor command also checks motion-prior robustness against static tennis-ball-like distractors. The strict command must fail until real licensed tennis data and iPhone 7+ benchmark results are supplied.

Validate real mobile evidence before strict release:

python scripts/real_world_readiness.py
python scripts/real_world_readiness.py --init-templates
python scripts/validate_mobile_benchmark.py --input runs/real_tennis/mobile_benchmark.json --require-iphone7
python scripts/validate_coreml_device_report.py --input runs/real_tennis/coreml_device_report.json

real_world_readiness.py reports which real-world evidence is present and which items still need human input. Use --init-templates to create data/real_tennis/manifest.json, runs/real_tennis/mobile_benchmark.json, and runs/real_tennis/coreml_device_report.json from the checked-in examples.

Known limitation: the current default model passes the clean synthetic motion set. Static tennis-ball-like distractors require motion-prior post-processing and still need real-court validation. See docs/completion_audit.md before treating the model as production-ready.

Mobile Target

Default input is 256x144, single-class output, and stride 4. The base model uses depthwise separable blocks and keeps the parameter budget under roughly one million parameters for the default width multiplier. For iPhone 7 and later, start with:

• input_width=256, input_height=144
• width_mult=0.50
• FP16 Core ML or int8 post-training quantization after calibration; TorchScript is also supported as a dependency-light mobile artifact
• 30 FPS video with motion ROI crop when available

See docs/research_notes.md for the checked 2024-2026 paper map and the clean-room/IP position. See docs/ip_cleanroom.md for the clean-room IP register and automated source scan. See docs/mobile_deployment.md for PyTorch Mobile Lite usage and the iPhone 7+ benchmark acceptance table. See docs/real_data_pipeline.md for real video extraction, labeling, validation, training, and device-test steps. See docs/real_video_handoff.md for the current tennis1.mp4 extraction outputs and labeling handoff. See docs/production_readiness.md for the current go/no-go status and required production gates. See docs/交付总结.md for a concise Chinese handoff summary. Use scripts/extract_frames.py to bootstrap annotated JSON from a licensed rally video before labeling boxes. Use scripts/annotations_csv.py to export/import a simple CSV labeling sheet. Use scripts/visualize_annotations.py to render bbox previews before training.