umadetector

YOLO object detection dataset and training pipeline for Umamusume Pretty Derby.

---

Setup

pip install -r requirements.txt

Requires xdotool on Linux for screen capture:

sudo pacman -S xdotool   # Arch/CachyOS

---

Workflow

1. Capture screenshots

Manual (F1 hotkey):

python screenshot_capture.py

Press F1 to save a screenshot, Esc to stop.

Autoplay loop:

python portrait_training_data_capture.py

Saves one screenshot per second while autoplay is running. Ctrl+C to stop.

Screenshots are saved to dataset/images/raw/ at 1075×1440, matching the training dataset format. No further cropping needed.

---

2. Label images

Open raw images in xanylabeling and draw bounding boxes using the class names in dataset/classes.txt.

To pre-label with an existing model (then correct manually):

python autolabel.py              # label all unlabelled images
python autolabel.py --conf 0.5   # stricter confidence
python autolabel.py --preview    # save annotated previews to test_output/autolabel/
python autolabel.py --overwrite  # redo already-labelled images

Autolabel writes xanylabeling-compatible .json files alongside the images in dataset/images/raw/. Images with no detections are skipped and left unlabelled for manual review.

---

3. Convert JSON → YOLO format

python yolo_convert.py              # convert all splits
python yolo_convert.py --split race # convert one split only

Reads dataset/proposed_split.txt to determine which classes belong to each split. Outputs per-split YOLO .txt label files to dataset/labels/{split}/raw/, with class IDs re-indexed from 0 for each split.

---

4. Split into train/val

python split_dataset.py

Splits labelled images 90/10 into per-split dataset/splits/{split}/images/train|val directories using hard links (so YOLO label resolution works correctly). Also generates a data.yaml for each split. Only images that have at least one label in any split are included. The shuffle is seeded so adding new unlabelled images doesn't reshuffle existing splits.

---

5. Train

python train.py                      # train all splits sequentially
python train.py --split training     # train one split only

Trains yolo26n for 100 epochs per split with early stopping (patience=20). Best weights saved to runs/uma_detector_{split}/weights/best.pt.

---

6. Test the model

python test_model.py                          # snapshot from game window
python test_model.py path/to/image.jpg        # run on a saved image
python test_model.py --conf 0.4               # override confidence threshold
python test_model.py --split training         # run only one split model

Crops the capture to the game play area and resizes to training dimensions before inference, so detection works correctly at any window size or resolution. Annotated result saved to test_output/.

---

Dataset structure

dataset/
  classes.txt               — full class list (for xanylabeling --labels)
  proposed_split.txt        — defines which classes belong to each model split
  images/
    raw/                    — all captured screenshots (1075×1440) + xanylabeling JSONs
  labels/
    {split}/
      raw/                  — per-split YOLO .txt label files (from yolo_convert.py)
  splits/
    {split}/
      data.yaml             — YOLO training config for this split
      images/
        train/              — hard-linked images for training (90%)
        val/                — hard-linked images for validation (10%)
      labels/
        train/              — label files for training images
        val/                — label files for val images

---

Model splits

Classes are divided across 5 models defined in dataset/proposed_split.txt:

Split	Classes	Description
ui_detection	30	General HUD elements, buttons, skill cards, aptitudes, race results, plus/minus
training	16	Training screen — tiles, cards, support members
event	4	Event choices and banners
race	4	Race screen elements
dreams	14	Dreams mode — team members, stats, DP

---

todo

Dataset wise:
diversify support cards (friendship_card especially — label varied card art)
skill hints
aptitudes, skill_card_na (no annotations yet)
red desire + espoir city portraits
inheritance
race specific elements