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

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