|
|
--- |
|
|
license: apache-2.0 |
|
|
tags: |
|
|
- computer-vision |
|
|
- chess |
|
|
- yolo |
|
|
- segmentation |
|
|
- instance-segmentation |
|
|
|
|
|
model-index: |
|
|
- name: chess_board_segmentation |
|
|
results: [] |
|
|
--- |
|
|
|
|
|
# chess_board_segmentation |
|
|
|
|
|
ChessBoard segmentation model trained with YOLO |
|
|
|
|
|
## Model Details |
|
|
|
|
|
- **Model Type**: YOLO Segmentation |
|
|
- **Task**: instance-segmentation |
|
|
- **License**: apache-2.0 |
|
|
|
|
|
## Usage |
|
|
|
|
|
### Loading the Model |
|
|
|
|
|
```python |
|
|
from src.chess_board_detection.yolo.segmentation.segmentation_model import ChessBoardSegmentationModel |
|
|
|
|
|
# Load the model |
|
|
model = ChessBoardSegmentationModel(model_path="path/to/downloaded/model.pt") |
|
|
|
|
|
# Get polygon coordinates for a chessboard |
|
|
polygon_info, is_valid = model.get_polygon_coordinates("path/to/chessboard_image.jpg") |
|
|
|
|
|
if is_valid: |
|
|
print(f"Detected chessboard polygon: {polygon_info}") |
|
|
|
|
|
# Extract corners from the segmentation |
|
|
corners = model.extract_corners_from_segmentation( |
|
|
"path/to/chessboard_image.jpg", |
|
|
polygon_info |
|
|
) |
|
|
print(f"Extracted corners: {corners}") |
|
|
|
|
|
# Visualize results |
|
|
model.plot_eval("path/to/chessboard_image.jpg", show=True) |
|
|
``` |
|
|
|
|
|
### Direct YOLO Usage |
|
|
|
|
|
```python |
|
|
from ultralytics import YOLO |
|
|
|
|
|
# Load the model |
|
|
model = YOLO("path/to/downloaded/model.pt") |
|
|
|
|
|
# Run segmentation |
|
|
results = model("path/to/chessboard_image.jpg") |
|
|
|
|
|
# Get masks and polygons |
|
|
for result in results: |
|
|
if result.masks is not None: |
|
|
for mask in result.masks: |
|
|
polygon = mask.xy[0] # Polygon coordinates |
|
|
print(f"Polygon points: {polygon}") |
|
|
``` |
|
|
|
|
|
### Training Data Format |
|
|
|
|
|
This model expects YOLO segmentation format with polygon annotations: |
|
|
|
|
|
```yaml |
|
|
# data.yaml |
|
|
train: path/to/train/images |
|
|
val: path/to/val/images |
|
|
nc: 1 |
|
|
names: ['chessboard'] |
|
|
``` |
|
|
|
|
|
With corresponding label files containing polygon coordinates: |
|
|
``` |
|
|
# labels/image.txt |
|
|
0 x1 y1 x2 y2 x3 y3 x4 y4 ... # normalized coordinates |
|
|
``` |
|
|
|
|
|
## Training |
|
|
|
|
|
This model was trained using the ChessBoard Segmentation training pipeline: |
|
|
|
|
|
```bash |
|
|
python src/chess_board_detection/yolo/segmentation/train_segmentation.py \ |
|
|
--data data/chessboard_segmentation/chess-board-3/data.yaml \ |
|
|
--epochs 100 \ |
|
|
--batch 16 \ |
|
|
--pretrained-model yolov8s-seg.pt |
|
|
``` |
|
|
|
|
|
## Model Performance |
|
|
|
|
|
<!-- Add performance metrics here after training --> |
|
|
|
|
|
## Citation |
|
|
|
|
|
If you use this model in your research, please cite: |
|
|
|
|
|
```bibtex |
|
|
@misc{dopaul_chess_board_segmentation, |
|
|
title={chess_board_segmentation}, |
|
|
author={dopaul}, |
|
|
year={2024}, |
|
|
publisher={Hugging Face}, |
|
|
url={https://huggingface.co/dopaul/chess_board_segmentation} |
|
|
} |
|
|
``` |
|
|
|
