1use std::fmt;
2
3use image::RgbImage;
4
5use crate::error::Result;
6use crate::error::XqVisionError;
7
8pub const BOARD_RANKS: usize = 10;
9pub const BOARD_FILES: usize = 9;
10pub const BOARD_CELLS: usize = BOARD_RANKS * BOARD_FILES;
11pub const PIECE_CLASSES: usize = 16;
12
13#[derive(Debug, Clone, Copy, PartialEq)]
14pub struct Point2f {
15 pub x: f32,
16 pub y: f32,
17}
18
19impl Point2f {
20 #[must_use]
21 pub const fn new(x: f32, y: f32) -> Self { Self { x, y } }
22
23 #[must_use]
24 pub const fn to_array(self) -> [f32; 2] { [self.x, self.y] }
25}
26
27impl From<[f32; 2]> for Point2f {
28 fn from(value: [f32; 2]) -> Self { Self::new(value[0], value[1]) }
29}
30
31impl From<Point2f> for [f32; 2] {
32 fn from(value: Point2f) -> Self { value.to_array() }
33}
34
35#[derive(Debug, Clone, Copy, PartialEq)]
36pub struct RectF {
37 pub x_min: f32,
38 pub y_min: f32,
39 pub x_max: f32,
40 pub y_max: f32,
41}
42
43impl RectF {
44 #[must_use]
45 pub const fn new(x_min: f32, y_min: f32, x_max: f32, y_max: f32) -> Self { Self { x_min, y_min, x_max, y_max } }
46
47 #[must_use]
48 pub fn from_image(image: &RgbImage) -> Self { Self::new(0.0, 0.0, image.width() as f32, image.height() as f32) }
49
50 #[must_use]
51 pub fn width(self) -> f32 { self.x_max - self.x_min }
52
53 #[must_use]
54 pub fn height(self) -> f32 { self.y_max - self.y_min }
55
56 pub(crate) fn validate(self) -> Result<()> {
57 if self.width() <= 0.0 || self.height() <= 0.0 {
58 return Err(XqVisionError::InvalidGeometry("bounding box must have positive size"));
59 }
60 Ok(())
61 }
62}
63
64#[derive(Debug, Clone, Copy, PartialEq)]
65pub struct BoardCorners {
66 pub top_left: Point2f,
67 pub top_right: Point2f,
68 pub bottom_left: Point2f,
69 pub bottom_right: Point2f,
70}
71
72impl BoardCorners {
73 #[must_use]
74 pub const fn new(top_left: Point2f, top_right: Point2f, bottom_left: Point2f, bottom_right: Point2f) -> Self {
75 Self { top_left, top_right, bottom_left, bottom_right }
76 }
77
78 #[must_use]
79 pub fn as_points(self) -> [Point2f; 4] { [self.top_left, self.top_right, self.bottom_left, self.bottom_right] }
80
81 #[must_use]
82 pub fn as_arrays(self) -> [[f32; 2]; 4] { self.as_points().map(Point2f::to_array) }
83
84 pub(crate) fn validate(self) -> Result<()> {
85 let points = self.as_points();
86 if points.iter().any(|point| !point.x.is_finite() || !point.y.is_finite()) {
87 return Err(XqVisionError::InvalidGeometry("board corners must be finite"));
88 }
89 Ok(())
90 }
91}
92
93impl From<[[f32; 2]; 4]> for BoardCorners {
94 fn from(value: [[f32; 2]; 4]) -> Self {
95 Self::new(value[0].into(), value[1].into(), value[2].into(), value[3].into())
96 }
97}
98
99#[derive(Debug, Clone)]
100pub struct BoardImage {
101 image: RgbImage,
102}
103
104impl BoardImage {
105 #[must_use]
106 pub fn new(image: RgbImage) -> Self { Self { image } }
107
108 #[must_use]
109 pub fn as_image(&self) -> &RgbImage { &self.image }
110
111 #[must_use]
112 pub fn into_image(self) -> RgbImage { self.image }
113
114 #[must_use]
115 pub fn width(&self) -> u32 { self.image.width() }
116
117 #[must_use]
118 pub fn height(&self) -> u32 { self.image.height() }
119}
120
121#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
122pub struct BoardCoord {
123 pub rank: usize,
124 pub file: usize,
125}
126
127impl BoardCoord {
128 #[must_use]
129 pub const fn new(rank: usize, file: usize) -> Self { Self { rank, file } }
130}
131
132#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
133pub enum Side {
134 Red,
135 Black,
136}
137
138impl Side {
139 #[must_use]
140 pub const fn fen_char(self) -> char {
141 match self {
142 Self::Red => 'w',
143 Self::Black => 'b',
144 }
145 }
146}
147
148#[repr(u8)]
149#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
150pub enum PieceKind {
151 Empty = 0,
152 Unknown = 1,
153 RedKing = 2,
154 RedAdvisor = 3,
155 RedBishop = 4,
156 RedKnight = 5,
157 RedRook = 6,
158 RedCannon = 7,
159 RedPawn = 8,
160 BlackKing = 9,
161 BlackAdvisor = 10,
162 BlackBishop = 11,
163 BlackKnight = 12,
164 BlackRook = 13,
165 BlackCannon = 14,
166 BlackPawn = 15,
167}
168
169impl PieceKind {
170 pub const ALL: [Self; PIECE_CLASSES] = [
171 Self::Empty,
172 Self::Unknown,
173 Self::RedKing,
174 Self::RedAdvisor,
175 Self::RedBishop,
176 Self::RedKnight,
177 Self::RedRook,
178 Self::RedCannon,
179 Self::RedPawn,
180 Self::BlackKing,
181 Self::BlackAdvisor,
182 Self::BlackBishop,
183 Self::BlackKnight,
184 Self::BlackRook,
185 Self::BlackCannon,
186 Self::BlackPawn,
187 ];
188
189 #[must_use]
190 pub const fn index(self) -> u8 { self as u8 }
191
192 #[must_use]
193 pub const fn short(self) -> char {
194 match self {
195 Self::Empty => '.',
196 Self::Unknown => 'x',
197 Self::RedKing => 'K',
198 Self::RedAdvisor => 'A',
199 Self::RedBishop => 'B',
200 Self::RedKnight => 'N',
201 Self::RedRook => 'R',
202 Self::RedCannon => 'C',
203 Self::RedPawn => 'P',
204 Self::BlackKing => 'k',
205 Self::BlackAdvisor => 'a',
206 Self::BlackBishop => 'b',
207 Self::BlackKnight => 'n',
208 Self::BlackRook => 'r',
209 Self::BlackCannon => 'c',
210 Self::BlackPawn => 'p',
211 }
212 }
213
214 pub fn from_index(index: u8) -> Result<Self> { Self::try_from(index) }
215}
216
217impl TryFrom<u8> for PieceKind {
218 type Error = XqVisionError;
219
220 fn try_from(value: u8) -> Result<Self> {
221 match value {
222 0 => Ok(Self::Empty),
223 1 => Ok(Self::Unknown),
224 2 => Ok(Self::RedKing),
225 3 => Ok(Self::RedAdvisor),
226 4 => Ok(Self::RedBishop),
227 5 => Ok(Self::RedKnight),
228 6 => Ok(Self::RedRook),
229 7 => Ok(Self::RedCannon),
230 8 => Ok(Self::RedPawn),
231 9 => Ok(Self::BlackKing),
232 10 => Ok(Self::BlackAdvisor),
233 11 => Ok(Self::BlackBishop),
234 12 => Ok(Self::BlackKnight),
235 13 => Ok(Self::BlackRook),
236 14 => Ok(Self::BlackCannon),
237 15 => Ok(Self::BlackPawn),
238 other => Err(XqVisionError::InvalidPieceIndex(other)),
239 }
240 }
241}
242
243impl fmt::Display for PieceKind {
244 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(&self.short().to_string()) }
245}
246
247#[derive(Debug, Clone, Copy, PartialEq)]
248pub struct CellPrediction {
249 pub coord: BoardCoord,
250 pub piece: PieceKind,
251 pub confidence: f32,
252}
253
254impl CellPrediction {
255 #[must_use]
256 pub const fn new(coord: BoardCoord, piece: PieceKind, confidence: f32) -> Self { Self { coord, piece, confidence } }
257}
258
259#[derive(Debug, Clone)]
260pub struct RecognitionResult {
261 corners: BoardCorners,
262 corner_scores: [f32; 4],
263 board: BoardImage,
264 pieces: crate::pieces::PieceRecognition,
265 user_side: Side,
266}
267
268impl RecognitionResult {
269 #[must_use]
270 pub(crate) fn new(
271 corners: BoardCorners, corner_scores: [f32; 4], board: BoardImage, pieces: crate::pieces::PieceRecognition,
272 user_side: Side,
273 ) -> Self {
274 Self { corners, corner_scores, board, pieces, user_side }
275 }
276
277 #[must_use]
278 pub fn corners(&self) -> BoardCorners { self.corners }
279
280 #[must_use]
281 pub fn corner_scores(&self) -> [f32; 4] { self.corner_scores }
282
283 #[must_use]
284 pub fn board(&self) -> &BoardImage { &self.board }
285
286 #[must_use]
287 pub fn pieces(&self) -> &crate::pieces::PieceRecognition { &self.pieces }
288
289 #[must_use]
290 pub fn user_side(&self) -> Side { self.user_side }
291
292 #[must_use]
293 pub fn to_fen(&self) -> String { self.pieces.to_fen(self.user_side) }
294
295 #[must_use]
296 pub fn into_parts(self) -> (BoardCorners, [f32; 4], BoardImage, crate::pieces::PieceRecognition, Side) {
297 (self.corners, self.corner_scores, self.board, self.pieces, self.user_side)
298 }
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304
305 #[test]
306 fn piece_kind_round_trips_all_indices() -> Result<()> {
307 for expected in PieceKind::ALL {
308 assert_eq!(PieceKind::from_index(expected.index())?, expected);
309 }
310 Ok(())
311 }
312
313 #[test]
314 fn invalid_piece_index_returns_error() {
315 assert!(matches!(PieceKind::from_index(99), Err(XqVisionError::InvalidPieceIndex(99))));
316 }
317}