<div align="center">

# Chess CV

<p align="center">
  <a href="https://github.com/S1M0N38/chess-cv/releases">
    <img alt="GitHub release" src="https://img.shields.io/github/v/release/S1M0N38/chess-cv?include_prereleases&sort=semver&style=for-the-badge&logo=github"/>
  </a>
  <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-yellow.svg?style=for-the-badge"/>
  </a>
  <a href="https://s1m0n38.github.io/chess-cv/llms-full.txt">
    <img alt="llms.txt" src="https://img.shields.io/badge/llms.txt-grey?style=for-the-badge"/>
  </a>
</p>

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

*CNN-based chess piece classifier*

</div>

---

A machine learning project that trains lightweight CNNs (156k parameters each) from scratch to classify chess pieces and arrow annotations from 32×32 pixel square images. The project includes two 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

Both 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')

# Make predictions
piece_predictions = pieces_model(image_tensor)
arrow_predictions = arrows_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 two 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.95%  |      99.95%      |
| [S1M0N38/chess-cv-openboard](https://huggingface.co/datasets/S1M0N38/chess-cv-openboard) \*     |    -     |      98.91%      |
| [S1M0N38/chess-cv-chessvision](https://huggingface.co/datasets/S1M0N38/chess-cv-chessvision) \* |    -     |      94.05%      |

\* *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

## 📚 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
- **[Train and Evaluate](https://s1m0n38.github.io/chess-cv/train-and-eval/)** – Training, evaluation, and deployment
- **[Inference](https://s1m0n38.github.io/chess-cv/inference/)** – Using pre-trained models and the library
- **[Architecture](https://s1m0n38.github.io/chess-cv/architecture/)** – Model design and performance characteristics

## 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>
