<div align="center">

# Chess CV

<p align="center">
  <a href="https://pypi.org/project/chess-cv/">
    <img alt="PyPI" src="https://img.shields.io/pypi/v/chess-cv?style=for-the-badge&logo=pypi&logoColor=white"/>
  </a>
  <a href="https://www.python.org/downloads/">
    <img alt="Python 3.13+" src="https://img.shields.io/badge/python-3.13+-blue.svg?style=for-the-badge&logo=python&logoColor=white"/>
  </a>
  <a href="https://opensource.org/licenses/MIT">
    <img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-green.svg?style=for-the-badge&color=0d990d"/>
  </a>
  <a href="https://huggingface.co/S1M0N38/chess-cv">
    <img alt="Hugging Face" src="https://img.shields.io/badge/Hugging%20Face-FFD21E?style=for-the-badge&label=model&logoColor=white&labelColor=grey"/>
  </a>
</p>

<img src="docs/assets/model.svg" alt="Model Architecture" width="600">

*Lightweight CNNs for chess board analysis*

</div>

---

A machine learning project that trains lightweight CNNs (156k parameters each) from scratch to classify chess pieces, arrow annotations, and piece centering from 32×32 pixel square images. The project includes three specialized models:

- **Pieces Model**: Classifies 13 classes (6 white pieces, 6 black pieces, empty squares) for board state recognition and FEN generation
- **Arrows Model**: Classifies 49 classes representing arrow overlay patterns for detecting and reconstructing chess analysis annotations
- **Snap Model**: Classifies 2 classes (centered vs off-centered pieces) for automated board analysis and piece positioning validation

All models are trained on synthetic data generated by combining 55 board styles from chess.com and lichess with piece sets and arrow overlays, achieving robust recognition across various visual styles.

## ⚡️ Quick Start

```bash
pip install chess-cv
```

Then use pre-trained models with bundled weights:

```python
from chess_cv import load_bundled_model

# Load pre-trained models (weights included in package)
pieces_model = load_bundled_model('pieces')
arrows_model = load_bundled_model('arrows')
snap_model = load_bundled_model('snap')

# Make predictions
piece_predictions = pieces_model(image_tensor)
arrow_predictions = arrows_model(image_tensor)
snap_predictions = snap_model(image_tensor)
```

For advanced usage or custom weight paths:

```python
from chess_cv import get_bundled_weight_path
from chess_cv.model import SimpleCNN
import mlx.core as mx

# Get path to bundled weights
weight_path = get_bundled_weight_path('pieces')

# Or download from HuggingFace Hub for latest version
# from huggingface_hub import hf_hub_download
# weight_path = hf_hub_download(repo_id="S1M0N38/chess-cv", filename="pieces.safetensors")

# Load model manually
model = SimpleCNN(num_classes=13)
weights = mx.load(str(weight_path))
model.load_weights(list(weights.items()))
```

## ✨ Features

**Lightweight Architecture**

- 156k parameter CNN optimized for 32×32px images
- MLX framework for efficient training
- Aggressive data augmentation for robust generalization

**Complete Pipeline**

- Synthetic data generation from board/piece combinations
- Training with early stopping and checkpointing
- Comprehensive evaluation with confusion matrices
- Optional Weights & Biases integration for experiment tracking
- Hugging Face Hub deployment for model sharing

## 🎯 Models

This project includes three specialized models for chess board analysis:

### ♟️ Pieces Model (`pieces.safetensors`)

Classifies chess square images into **13 classes**: 6 white pieces (wP, wN, wB, wR, wQ, wK), 6 black pieces (bP, bN, bB, bR, bQ, bK), and empty squares (xx). Designed for board state recognition and FEN generation.

**Training:** ~93,000 synthetic images with aggressive augmentation (arrow overlays, flips, rotation, color jitter)

**Performance:**

| Dataset                                                                                         | Accuracy | F1-Score (Macro) |
| ----------------------------------------------------------------------------------------------- | :------: | :--------------: |
| Test Data                                                                                       |  99.93%  |      99.93%      |
| [S1M0N38/chess-cv-openboard](https://huggingface.co/datasets/S1M0N38/chess-cv-openboard) \*     |    -     |      98.84%      |
| [S1M0N38/chess-cv-chessvision](https://huggingface.co/datasets/S1M0N38/chess-cv-chessvision) \* |    -     |      94.33%      |

\* *Dataset with unbalanced class distribution (e.g. many more samples for empty square class), so accuracy is not representative.*

### ↗ Arrows Model (`arrows.safetensors`)

Classifies chess square images into **49 classes** representing arrow overlay patterns: 20 arrow heads, 12 tails, 8 middle segments, 4 corners, and empty squares. Enables detection and reconstruction of arrow annotations in chess interfaces. The NSEW naming (North/South/East/West) indicates arrow orientation, with corners handling knight-move arrow patterns.

**Training:** ~3.14M synthetic images (20 epochs, batch size 128) with conservative augmentation (no flips/rotation to preserve arrow directionality)

**Performance:**

| Dataset               | Accuracy | F1-Score (Macro) |
| --------------------- | :------: | :--------------: |
| Test Data (synthetic) |  99.97%  |      99.97%      |

**Limitation:** Single arrow component per square only

### 📐 Snap Model (`snap.safetensors`)

Classifies chess square images into **2 classes**: centered ("ok") and off-centered ("bad") pieces. Designed for automated board analysis and piece positioning validation, helping ensure proper piece placement in digital chess interfaces.

**Training:** ~1.4M synthetic images (centered and off-centered piece positions) with conservative augmentation (no rotation to preserve centering semantics)

**Performance:**

| Dataset               | Accuracy | F1-Score (Macro) |
| --------------------- | :------: | :--------------: |
| Test Data (synthetic) |  99.92%  |      99.92%      |

**Use Cases:** Automated board state validation, piece positioning quality control, and chess interface usability testing.

For detailed information about the snap model architecture, training strategies, and data augmentation pipeline, see the [Architecture](https://s1m0n38.github.io/chess-cv/architecture/) and [Data Augmentations](https://s1m0n38.github.io/chess-cv/data-augmentations/) documentation.

## 📚 Documentation

For detailed documentation, visit [s1m0n38.github.io/chess-cv](https://s1m0n38.github.io/chess-cv/) or explore:

- **[Setup Guide](https://s1m0n38.github.io/chess-cv/setup/)** – Installation and data preparation
- **[Inference](https://s1m0n38.github.io/chess-cv/inference/)** – Using pre-trained models and the library
- **[Train and Evaluate](https://s1m0n38.github.io/chess-cv/train-and-eval/)** – Training, evaluation, and deployment
- **[Architecture](https://s1m0n38.github.io/chess-cv/architecture/)** – Model design and performance characteristics
- **[Data Augmentations](https://s1m0n38.github.io/chess-cv/data-augmentations/)** – Augmentation strategies and pipelines

## License

This project is licensed under the MIT License – see the [LICENSE](LICENSE) file for details.

---

<div align="center">

[Get Started](#quick-start) • [Contribute](CONTRIBUTING.md) • [Report Issues](https://github.com/S1M0N38/chess-cv/issues)

</div>
