balatro-yolo

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

1. gameplay video frames are loaded

2. frames are optionally downsampled via frame_stride

3. batches of frames are passed through YOLOv8

4. 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