Balatro YOLOv8 Detection

Fine-tuned YOLOv8 model for detecting gameplay objects in Balatro from video frames.

This model is trained on large-scale annotated gameplay data and is designed to detect all relevant in-game objects required for reconstructing structured game state. It serves as the first stage in a full computer vision pipeline used for dataset generation and imitation learning.

For information on how this model was trained, see Balatro CV Pipeline repository

Repository

balatro-yolo/
├── assets/
│   ├── train_batch0.jpg
│   ├── val_batch0_pred.jpg
│   └── results.png
├── weights/
│   └── model.pt
├── config.json
├── inference.py
├── data.yaml
└── README.md

Detected Objects

The model predicts bounding boxes and class labels for a wide range of gameplay elements, including:

cards, jokers, consumables, packs, tags, vouchers
editions, modifiers, seals, stickers
boss blind tokens
white, black and gold stake tokens
facedown and debuffed cards
deck types

This enables full reconstruction of game state when combined with tracking and OCR.

Example Pipeline

video frame
   ↓
YOLOv8 detection (this model)
   ↓
bounding boxes + class IDs + confidence
   ↓
object tracking
   ↓
state composition
   ↓
structured dataset

How It Works

gameplay video frames are loaded
frames are optionally downsampled via frame_stride
batches of frames are passed through YOLOv8
detections are written to CSV with:
- frame index
- timestamp
- class ID
- confidence
- bounding box coordinates

The output is used downstream for:

object tracking
metadata composition (editions, stickers, etc.)
event detection
imitation learning

Installation

pip install ultralytics opencv-python torch numpy

Usage

Single-frame / image inference

from ultralytics import YOLO

model = YOLO("weights/model.pt")
results = model("frame.png")

for r in results:
    boxes = r.boxes.xyxy
    classes = r.boxes.cls
    confs = r.boxes.conf

Video visualization

import cv2
from ultralytics import YOLO

model = YOLO("weights/model.pt")

cap = cv2.VideoCapture("clip.mp4")

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = model.predict(source=frame, imgsz=640, half=True, verbose=False)
    annotated = results[0].plot()

    cv2.imshow("output", annotated)
    if cv2.waitKey(1) == 27:
        break

cap.release()
cv2.destroyAllWindows()

Bulk processing (multi-GPU)

nohup python scripts/detect/detect_bulk.py \
  --project-root $PROJECT_ROOT \
  --video-dir $VIDEO_DIR \
  --model-path $MODEL_PATH \
  --gpus 0,1,2,3,4,5,6,7 \
  --frame-stride 2 \
  --imgsz 640 \
  --batch-size 32 \
  --recursive \
  --half \
  > $LOG_FILE 2>&1 &

Features

multi-GPU parallel processing
batched inference
frame skipping via frame_stride
automatic CSV export per video

Output Format

Each detection is written to CSV with the following schema:

video_id
frame_idx
timestamp_ms
class_id
conf
x1
y1
x2
y2
model_name
model_version
box_layout_version

This format is designed for downstream processing into:

parquet datasets
object tracking
composed object state

Training

The model was trained using YOLOv8 with a custom dataset:

nc: 400 classes
imgsz: 640
epochs: 150
batch: 56

Training command

from ultralytics import YOLO

model = YOLO("yolov8m.pt")
model.train(
    data="data.yaml",
    epochs=150,
    patience=40,
    imgsz=640,
    batch=56,
    workers=16,
    device="0,1,2,3,4,5,6,7"
)

Example Predictions

Training Results

Final performance

precision: 0.999
recall:    0.996
mAP50:     0.995
mAP50-95:  0.993

Dataset

The model is trained on a large-scale annotated dataset of Balatro gameplay frames.

~13k images
~1M object instances
400 classes

Project Context

This model is part of a full computer vision system for reconstructing gameplay state from video.

video
   ↓
object detection (this model)
   ↓
object tracking
   ↓
composed object state
   ↓
observations dataset
   ↓
ML agent training

The final goal is to enable large-scale imitation learning without requiring access to the game engine.

Notes

optimized specifically for Balatro gameplay
not intended for general object detection tasks
designed to integrate with tracking and OCR pipelines
supports large-scale offline dataset generation

License

MIT

Downloads last month: 2