Python utilities for training YOLO ball detection models and working with pool table datasets.

# Create virtual environment (from project root)
python3 -m venv .venv
source .venv/bin/activate

# Install dependencies
pip install -r requirements.txt

The training pipeline:

1. Prepare dataset with images and YOLO-format labels
2. Remap class labels to standard format (0=black, 1=cue, 2=solid, 3=stripe)
3. Split into train/val/test sets
4. Train with Ultralytics YOLO
5. Export to ONNX for mobile deployment

1. Edit the CONFIG dict in model_table.py to point to your local dataset:

   CONFIG = {
       "srcImgDir": "data/my_dataset/images",  # Path to your images
       "srcLblDir": "data/my_dataset/labels",  # Path to your YOLO labels
       ...
   }

2. Run training:

   cd python
   python model_table.py

Create a dataset directory with the following structure:

data/my_dataset/
├── images/
│   ├── image001.jpg
│   ├── image002.jpg
│   └── ...
└── labels/
    ├── image001.txt
    ├── image002.txt
    └── ...

Label format (YOLO): Each .txt file contains one line per object:

<class_id> <center_x> <center_y> <width> <height>

All coordinates are normalized (0-1). Example:

0 0.5123 0.3456 0.0234 0.0312
2 0.7891 0.6543 0.0198 0.0287

Class IDs:

If your images show full table views at angles, transform them to normalized top-down perspective:

python transform_dataset.py data/my_dataset -o data/my_dataset_transformed

# Options:
#   --visualize          Show each transformation
#   --require-valid      Skip images with bad orientation detection

This:

• Detects the table quadrilateral using C++ FFI
• Applies perspective correction to 864x1680 (ShotStudio dimensions)
• Transforms ball bounding boxes accordingly

Edit model_table.py or create a JSON config file:

CONFIG = {
    "srcImgDir": "data/my_dataset/images",
    "srcLblDir": "data/my_dataset/labels",
    "dstRoot": "/tmp/workdir",
    "oldToNewMap": {0: 0, 1: 1, 2: 2, 3: 3},  # Class remapping
    "classNames": ["black", "cue", "solid", "stripe"],
    "split": [0.8, 0.10, 0.10],  # train/val/test ratios
    "trainer": {
        "model": "yolov8n.pt",      # Base model
        "hyp": "data/hyps/hyp.custom.yaml",  # Hyperparameters
        "epochs": 40,
        "imgsz": 1280,              # Input image size
        "batch": 4,                 # Batch size (reduce if OOM)
        "device": "mps",            # "mps", "cuda:0", or "cpu"
        "workers": 8,
        "project": "tableizer",
        "name": "my_model",         # Output directory name
    },
}

python model_table.py

Training outputs are saved to tableizer/<name>/:

tableizer/my_model/
├── weights/
│   ├── best.pt          # Best model weights
│   └── last.pt          # Final epoch weights
├── args.yaml            # Training configuration
├── results.csv          # Training metrics
└── *.png                # Training curves and confusion matrix

For mobile deployment, export to ONNX format:

yolo export model=tableizer/expN/weights/best.pt format=onnx device=cpu imgsz=1280 simplify=True dynamic=False opset=17 half=False

Copy the model to the Flutter app:

cp tableizer/expN/weights/best.onnx ../app/assets/detection_model.onnx

Custom hyperparameters are defined in data/hyps/hyp.custom.yaml:

# Learning rate
lr0: 0.01           # Initial learning rate
lrf: 0.01           # Final learning rate (lr0 * lrf)

# Augmentation
hsv_h: 0.015        # Hue augmentation
hsv_s: 0.7          # Saturation augmentation
hsv_v: 0.4          # Value augmentation
degrees: 5.0        # Rotation (+/- degrees)
translate: 0.1      # Translation (+/- fraction)
scale: 0.25         # Scale (+/- gain)
flipud: 0.1         # Vertical flip probability
fliplr: 0.5         # Horizontal flip probability
mosaic: 0.45        # Mosaic augmentation probability

# Loss weights
box: 0.1            # Box loss gain
cls: 0.5            # Classification loss gain
dfl: 1.5            # Distribution focal loss gain

# Detection
conf: 0.4           # Confidence threshold
iou: 0.8            # IoU threshold
max_det: 20         # Maximum detections per image

Test table and ball detection on images:

python detect_table.py path/to/image.jpg

# Options:
#   --model PATH         Path to YOLO model
#   --visualize          Show detection overlay
#   --rotation DEGREES   Rotate input image

Run detection on perspective-corrected images:

python detect_transformed_table.py path/to/image.jpg

Validate YOLO label files:

python check_labels.py data/my_dataset/labels/

Compare PyTorch and ONNX model outputs:

python compare_models.py

Python bindings for the C++ native library:

from tableizer_ffi import detect_table_cpp, initialize_ball_detector

# Detect table quadrilateral
result = detect_table_cpp(image_bgr, rotation_degrees=0)
quad_points = result["quad_points"]
orientation = result["orientation"]

# Initialize ball detector with ONNX model
detector = initialize_ball_detector("path/to/model.onnx")

If training crashes with OOM errors:

• Reduce batch size (try 2 or 1)
• Reduce imgsz (try 640)
• Reduce workers
• Use cache: disk instead of cache: True

• More data: Combine multiple datasets
• Data augmentation: Adjust hyperparameters in hyp.custom.yaml
• Larger model: Try yolov8s.pt or yolov8m.pt instead of yolov8n.pt
• More epochs: Increase training epochs (monitor for overfitting)
• Lower confidence: Reduce conf threshold for more detections

• Ensure consistent labeling across all images
• Include variety: different table colors, lighting, angles
• Balance classes: similar counts of each ball type
• Use transform_dataset.py to normalize perspective if images are at angles