xq-vision 0.1.3

High-performance ONNX recognition pipeline for Chinese chessboard corners and pieces.
Documentation
use image::RgbImage;

use crate::board::BoardDetector;
use crate::board::BoardDetectorConfig;
use crate::board::warp_board;
use crate::config::GraphOptimization;
use crate::config::ModelSource;
use crate::config::SessionConfig;
use crate::error::Result;
use crate::error::XqVisionError;
use crate::pieces::PieceRecognizer;
use crate::pieces::PieceRecognizerConfig;
use crate::session::ProviderFailure;
use crate::types::RecognitionResult;

pub struct XqVision {
    board_detector: BoardDetector,
    piece_recognizer: PieceRecognizer,
}

impl XqVision {
    #[must_use]
    pub fn builder() -> XqVisionBuilder { XqVisionBuilder::default() }

    pub fn recognize(&mut self, image: &RgbImage) -> Result<RecognitionResult> {
        let board_detection = self.board_detector.detect(image)?;
        let board = warp_board(image, board_detection.corners)?;
        let pieces = self.piece_recognizer.recognize(&board)?;
        let user_side = pieces.user_side()?;
        Ok(RecognitionResult::new(board_detection.corners, board_detection.scores, board, pieces, user_side))
    }

    #[must_use]
    pub fn board_detector(&self) -> &BoardDetector { &self.board_detector }

    #[must_use]
    pub fn board_detector_mut(&mut self) -> &mut BoardDetector { &mut self.board_detector }

    #[must_use]
    pub fn piece_recognizer(&self) -> &PieceRecognizer { &self.piece_recognizer }

    #[must_use]
    pub fn piece_recognizer_mut(&mut self) -> &mut PieceRecognizer { &mut self.piece_recognizer }
}

#[derive(Debug, Default, Clone)]
pub struct XqVisionBuilder {
    board_model: Option<ModelSource>,
    piece_model: Option<ModelSource>,
    board_config: BoardDetectorConfig,
    piece_config: PieceRecognizerConfig,
    session_config: SessionConfig,
}

impl XqVisionBuilder {
    #[must_use]
    pub fn board_model(mut self, model: impl Into<ModelSource>) -> Self {
        self.board_model = Some(model.into());
        self
    }

    #[must_use]
    pub fn piece_model(mut self, model: impl Into<ModelSource>) -> Self {
        self.piece_model = Some(model.into());
        self
    }

    #[must_use]
    pub fn board_config(mut self, config: BoardDetectorConfig) -> Self {
        self.board_config = config;
        self
    }

    #[must_use]
    pub fn piece_config(mut self, config: PieceRecognizerConfig) -> Self {
        self.piece_config = config;
        self
    }

    #[must_use]
    pub fn session_config(&self) -> &SessionConfig { &self.session_config }

    #[must_use]
    pub fn with_session_config(mut self, config: SessionConfig) -> Self {
        self.session_config = config;
        self
    }

    #[must_use]
    pub fn provider_failure(mut self, failure: ProviderFailure) -> Self {
        self.session_config = self.session_config.with_provider_failure(failure);
        self
    }

    #[must_use]
    pub fn graph_optimization(mut self, level: GraphOptimization) -> Self {
        self.session_config = self.session_config.with_graph_optimization(level);
        self
    }

    #[must_use]
    pub fn intra_threads(mut self, threads: usize) -> Self {
        self.session_config = self.session_config.with_intra_threads(threads);
        self
    }

    #[must_use]
    pub fn inter_threads(mut self, threads: usize) -> Self {
        self.session_config = self.session_config.with_inter_threads(threads);
        self
    }

    #[must_use]
    pub fn parallel_execution(mut self, enabled: bool) -> Self {
        self.session_config = self.session_config.with_parallel_execution(enabled);
        self
    }

    pub fn build(self) -> Result<XqVision> {
        let board_model = self.board_model.ok_or(XqVisionError::MissingModel { role: "board" })?;
        let piece_model = self.piece_model.ok_or(XqVisionError::MissingModel { role: "piece" })?;
        let board_detector = BoardDetector::with_config(board_model, self.board_config, self.session_config.clone())?;
        let piece_recognizer = PieceRecognizer::with_config(piece_model, self.piece_config, self.session_config)?;
        Ok(XqVision { board_detector, piece_recognizer })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn builder_requires_board_model() {
        let result = XqVision::builder().piece_model(ModelSource::memory([1_u8])).build();
        assert!(matches!(result, Err(XqVisionError::MissingModel { role: "board" })));
    }

    #[test]
    fn builder_requires_piece_model() {
        let result = XqVision::builder().board_model(ModelSource::memory([1_u8])).build();
        assert!(matches!(result, Err(XqVisionError::MissingModel { role: "piece" })));
    }
}