This commit is contained in:
TheMaddax 2025-01-13 12:42:20 -06:00
parent a46c15fc6f
commit 60c2217348
20 changed files with 38287 additions and 1183 deletions

View file

@ -0,0 +1,5 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict/>
</plist>

View file

@ -95,11 +95,11 @@ final class BoardDetector {
throw BoardDetectionError.invalidBoardDimensions throw BoardDetectionError.invalidBoardDimensions
} }
var position = ChessPosition()
var recognizedPieces = 0
// Reset piece counts before starting new scan // Reset piece counts before starting new scan
pieceRecognizer.resetPieceCounts() pieceRecognizer.resetCounts()
var squares = [BoardPosition: SquareClassification]()
var pieceCount = 0
// Process each square // Process each square
for rank in 0..<8 { for rank in 0..<8 {
@ -107,33 +107,27 @@ final class BoardDetector {
let boardPosition = BoardPosition(file: file, rank: rank) let boardPosition = BoardPosition(file: file, rank: rank)
let squareImage = try extractSquare(from: image, in: boardRect, at: boardPosition) let squareImage = try extractSquare(from: image, in: boardRect, at: boardPosition)
// Set the current position being analyzed // Recognize square content
pieceRecognizer.currentPosition = boardPosition let classification = try await pieceRecognizer.recognizeSquare(from: squareImage, row: rank, col: file)
squares[boardPosition] = classification
if let piece = try await pieceRecognizer.recognizePiece(from: squareImage) { // Count pieces for logging
position[boardPosition] = piece if !classification.isEmpty {
recognizedPieces += 1 pieceCount += 1
} }
} }
} }
// Validate piece counts after all squares are processed // Generate FEN string and create position
try pieceRecognizer.validatePieceCounts() let fenGenerator = FenGenerator()
let fen = fenGenerator.generateFen(from: squares)
print("\nPieces recognized: \(recognizedPieces)") guard let position = ChessPosition(fen: fen) else {
print("\nRecognized position:") print("ERROR: Failed to create position from FEN")
for rank in (0...7).reversed() { throw BoardDetectionError.imageProcessingFailed
var rankStr = "\(rank + 1) "
for file in 0...7 {
if let piece = position[BoardPosition(file: file, rank: rank)] {
rankStr += "\(piece.color == .white ? "w" : "b")\(piece.type.fenSymbol.uppercased()) "
} else {
rankStr += ".. "
}
}
print(rankStr)
} }
print(" a b c d e f g h")
print("\nPieces recognized: \(pieceCount)")
print("\nGenerated FEN: \(fen)")
guard position.isValid else { guard position.isValid else {
print("\nERROR: Invalid chess position") print("\nERROR: Invalid chess position")

View file

@ -1,86 +0,0 @@
[
{
"metadataOutputVersion" : "3.0",
"outputSchema" : [
{
"isOptional" : "0",
"formattedType" : "String",
"type" : "String",
"name" : "target",
"shortDescription" : ""
},
{
"isOptional" : "0",
"keyType" : "String",
"formattedType" : "Dictionary (String → Double)",
"type" : "Dictionary",
"name" : "targetProbability",
"shortDescription" : ""
}
],
"modelParameters" : [
],
"author" : "Chris Haulmark",
"specificationVersion" : 8,
"isUpdatable" : "0",
"stateSchema" : [
],
"availability" : {
"macOS" : "14.0",
"tvOS" : "17.0",
"visionOS" : "1.0",
"watchOS" : "unavailable",
"iOS" : "17.0",
"macCatalyst" : "17.0"
},
"modelType" : {
"name" : "MLModelType_imageClassifier",
"structure" : [
{
"name" : "MLModelType_visionFeaturePrint"
},
{
"name" : "MLModelType_glmClassifier"
}
]
},
"inputSchema" : [
{
"height" : "360",
"colorspace" : "BGR",
"isOptional" : "0",
"width" : "360",
"isColor" : "1",
"formattedType" : "Image (Color 360 × 360)",
"hasSizeFlexibility" : "0",
"type" : "Image",
"shortDescription" : "",
"name" : "image"
}
],
"classLabels" : [
"black_bishop",
"black_king",
"black_knight",
"black_pawn",
"black_queen",
"black_rook",
"white_bishop",
"white_king",
"white_knight",
"white_pawn",
"white_queen",
"white_rook"
],
"generatedClassName" : "ChessPieceClassifier",
"userDefinedMetadata" : {
"com.apple.createml.version" : "15.3.0",
"com.apple.createml.app.tag" : "150.3",
"com.apple.coreml.model.preview.type" : "imageClassifier",
"com.apple.createml.app.version" : "6.1"
},
"method" : "predict"
}
]

View file

@ -4,7 +4,7 @@
<dict> <dict>
<key>com.apple.security.app-sandbox</key> <key>com.apple.security.app-sandbox</key>
<true/> <true/>
<key>com.apple.security.files.user-selected.read-only</key> <key>com.apple.security.files.user-selected.read-write</key>
<true/> <true/>
<key>com.apple.security.device.screen-capture</key> <key>com.apple.security.device.screen-capture</key>
<true/> <true/>
@ -14,5 +14,11 @@
<true/> <true/>
<key>com.apple.security.ml.coreml</key> <key>com.apple.security.ml.coreml</key>
<true/> <true/>
<key>com.apple.security.files.user-selected.executable</key>
<true/>
<key>com.apple.security.files.downloads.read-write</key>
<true/>
<key>com.apple.security.files.user-selected.read-write</key>
<true/>
</dict> </dict>
</plist> </plist>

View file

@ -41,10 +41,25 @@ struct ChessPiece: Equatable {
} }
/// Represents a position on the chess board /// Represents a position on the chess board
struct BoardPosition: Equatable { struct BoardPosition: Equatable, Hashable {
let file: Int // 0-7 for a-h let file: Int // 0-7 for a-h
let rank: Int // 0-7 for 1-8 let rank: Int // 0-7 for 1-8
/// Initialize from file and rank indices (0-7)
init(file: Int, rank: Int) {
guard file >= 0, file < 8, rank >= 0, rank < 8 else {
fatalError("Invalid board position: file \(file), rank \(rank)")
}
self.file = file
self.rank = rank
}
/// Hash function implementation for Hashable conformance
func hash(into hasher: inout Hasher) {
hasher.combine(file)
hasher.combine(rank)
}
/// Initialize from algebraic notation (e.g., "e4") /// Initialize from algebraic notation (e.g., "e4")
init?(algebraic: String) { init?(algebraic: String) {
guard algebraic.count == 2, guard algebraic.count == 2,

View file

@ -0,0 +1,54 @@
import Foundation
struct SquareClassification {
let pieceType: PieceType?
let pieceColor: PieceColor?
let isHighlighted: Bool
var isEmpty: Bool {
return pieceType == nil
}
static func empty(highlighted: Bool = false) -> SquareClassification {
return SquareClassification(pieceType: nil, pieceColor: nil, isHighlighted: highlighted)
}
init(pieceType: PieceType?, pieceColor: PieceColor?, isHighlighted: Bool = false) {
self.pieceType = pieceType
self.pieceColor = pieceColor
self.isHighlighted = isHighlighted
}
init?(label: String) {
// Match exact categories from trained model
switch label {
case "white_pawn":
self.init(pieceType: .pawn, pieceColor: .white)
case "white_knight":
self.init(pieceType: .knight, pieceColor: .white)
case "white_bishop":
self.init(pieceType: .bishop, pieceColor: .white)
case "white_rook":
self.init(pieceType: .rook, pieceColor: .white)
case "white_queen":
self.init(pieceType: .queen, pieceColor: .white)
case "white_king":
self.init(pieceType: .king, pieceColor: .white)
case "black_pawn":
self.init(pieceType: .pawn, pieceColor: .black)
case "black_knight":
self.init(pieceType: .knight, pieceColor: .black)
case "black_bishop":
self.init(pieceType: .bishop, pieceColor: .black)
case "black_rook":
self.init(pieceType: .rook, pieceColor: .black)
case "black_queen":
self.init(pieceType: .queen, pieceColor: .black)
case "black_king":
self.init(pieceType: .king, pieceColor: .black)
default:
// Any unrecognized label returns an empty square
self.init(pieceType: nil, pieceColor: nil, isHighlighted: false)
}
}
}

View file

@ -0,0 +1,72 @@
import Foundation
/// A class responsible for generating FEN strings from square classifications
final class FenGenerator {
/// Generate a FEN string from a set of square classifications
/// - Parameter squares: Dictionary mapping board positions to their classifications
/// - Returns: FEN string representing the position
func generateFen(from squares: [BoardPosition: SquareClassification]) -> String {
var fen = ""
var emptyCount = 0
// Process each rank from top to bottom (8 to 1)
for rank in (0...7).reversed() {
// Process each file from left to right (a to h)
for file in 0...7 {
let position = BoardPosition(file: file, rank: rank)
guard let square = squares[position] else {
// If square is missing, treat as empty
emptyCount += 1
continue
}
if square.isEmpty {
// Count consecutive empty squares
emptyCount += 1
} else if let color = square.pieceColor,
let type = square.pieceType {
// If we had empty squares before this piece, add the count
if emptyCount > 0 {
fen += String(emptyCount)
emptyCount = 0
}
// Add the piece symbol
let symbol = pieceSymbol(color: color, type: type)
fen += symbol
}
}
// Add any remaining empty squares at end of rank
if emptyCount > 0 {
fen += String(emptyCount)
emptyCount = 0
}
// Add rank separator (except for last rank)
if rank > 0 {
fen += "/"
}
}
return fen
}
/// Get the FEN symbol for a piece
/// - Parameters:
/// - color: Color of the piece
/// - type: Type of the piece
/// - Returns: FEN symbol (uppercase for white, lowercase for black)
private func pieceSymbol(color: PieceColor, type: PieceType) -> String {
let symbol: String
switch type {
case .king: symbol = "K"
case .queen: symbol = "Q"
case .rook: symbol = "R"
case .bishop: symbol = "B"
case .knight: symbol = "N"
case .pawn: symbol = "P"
}
return color == .white ? symbol : symbol.lowercased()
}
}

View file

@ -0,0 +1,49 @@
import Foundation
/// Represents a detected chess move
struct DetectedMove {
/// The source square where the piece moved from
let from: BoardPosition
/// The target square where the piece moved to
let to: BoardPosition
/// The piece that was moved
let piece: ChessPiece
}
/// A class responsible for detecting moves from highlighted squares
final class MoveDetector {
/// Detect a move from a set of square classifications
/// - Parameter squares: Dictionary mapping board positions to their classifications
/// - Returns: Detected move if one is found, nil otherwise
func detectMove(from squares: [BoardPosition: SquareClassification]) -> DetectedMove? {
var sourceSquare: BoardPosition?
var targetSquare: BoardPosition?
var movedPiece: ChessPiece?
// Find highlighted squares
for (position, classification) in squares {
guard classification.isHighlighted else { continue }
if classification.isEmpty {
// Empty highlighted square is the source
sourceSquare = position
} else if let color = classification.pieceColor,
let type = classification.pieceType {
// Piece on highlighted square is the target
targetSquare = position
movedPiece = ChessPiece(type: type, color: color)
}
}
// Return move if we found both squares and the piece
if let from = sourceSquare,
let to = targetSquare,
let piece = movedPiece {
return DetectedMove(from: from, to: to, piece: piece)
}
return nil
}
}

View file

@ -10,7 +10,6 @@ enum PieceRecognitionError: Error {
case recognitionFailed(String) case recognitionFailed(String)
case lowConfidence case lowConfidence
case invalidInput case invalidInput
case invalidPosition(String)
} }
/// A class responsible for recognizing chess pieces from images /// A class responsible for recognizing chess pieces from images
@ -23,9 +22,11 @@ final class PieceRecognizer {
/// Shared CIContext for image processing /// Shared CIContext for image processing
private static let ciContext = CIContext() private static let ciContext = CIContext()
/// Piece counts for validation /// Latest classification results
private var whitePieceCount: [PieceType: Int] = [:] private var classificationResults: [VNClassificationObservation]?
private var blackPieceCount: [PieceType: Int] = [:]
/// Confidence threshold
private let confidenceThreshold: Float = 0.75
// MARK: - Initialization // MARK: - Initialization
@ -34,7 +35,7 @@ final class PieceRecognizer {
let bundle = Bundle.main let bundle = Bundle.main
// Load model from bundle // Load model
guard let modelURL = bundle.url(forResource: "ChessPieceClassifier", withExtension: "mlmodelc") else { guard let modelURL = bundle.url(forResource: "ChessPieceClassifier", withExtension: "mlmodelc") else {
print("ERROR: Model not found in bundle at \(bundle.bundlePath)") print("ERROR: Model not found in bundle at \(bundle.bundlePath)")
throw PieceRecognitionError.modelLoadError throw PieceRecognitionError.modelLoadError
@ -50,109 +51,11 @@ final class PieceRecognizer {
print("ERROR: Failed to load model - \(error)") print("ERROR: Failed to load model - \(error)")
throw PieceRecognitionError.modelLoadError throw PieceRecognitionError.modelLoadError
} }
resetPieceCounts()
} }
// MARK: - Recognition Methods // MARK: - Recognition Methods
/// Reset piece counts for new position func recognizeSquare(from image: CGImage, row: Int = 0, col: Int = 0) async throws -> SquareClassification {
func resetPieceCounts() {
whitePieceCount = [
.king: 0,
.queen: 0,
.rook: 0,
.bishop: 0,
.knight: 0,
.pawn: 0
]
blackPieceCount = [
.king: 0,
.queen: 0,
.rook: 0,
.bishop: 0,
.knight: 0,
.pawn: 0
]
}
/// Validate piece counts and ensure the position is legal
func validatePieceCounts() throws {
var errors: [String] = []
// Check white pieces
if whitePieceCount[.king] != 1 {
let count = whitePieceCount[.king] ?? 0
errors.append("Invalid white king count: \(count)")
}
if let count = whitePieceCount[.queen], count > 1 {
errors.append("Too many white queens: \(count)")
}
if let count = whitePieceCount[.rook], count > 2 {
errors.append("Too many white rooks: \(count)")
}
if let count = whitePieceCount[.bishop], count > 2 {
errors.append("Too many white bishops: \(count)")
}
if let count = whitePieceCount[.knight], count > 2 {
errors.append("Too many white knights: \(count)")
}
if let count = whitePieceCount[.pawn], count > 8 {
errors.append("Too many white pawns: \(count)")
}
// Check black pieces
if blackPieceCount[.king] ?? 0 != 1 {
let count = blackPieceCount[.king] ?? 0
errors.append("Invalid black king count: \(count)")
}
if let count = blackPieceCount[.queen], count > 1 {
errors.append("Too many black queens: \(count)")
}
if let count = blackPieceCount[.rook], count > 2 {
errors.append("Too many black rooks: \(count)")
}
if let count = blackPieceCount[.bishop], count > 2 {
errors.append("Too many black bishops: \(count)")
}
if let count = blackPieceCount[.knight], count > 2 {
errors.append("Too many black knights: \(count)")
}
if let count = blackPieceCount[.pawn], count > 8 {
errors.append("Too many black pawns: \(count)")
}
if !errors.isEmpty {
throw PieceRecognitionError.invalidPosition(errors.joined(separator: ", "))
}
}
/// Update piece count
private func updatePieceCount(piece: ChessPiece) {
if piece.color == .white {
whitePieceCount[piece.type] = (whitePieceCount[piece.type] ?? 0) + 1
} else {
blackPieceCount[piece.type] = (blackPieceCount[piece.type] ?? 0) + 1
}
}
/// Check if adding this piece would exceed limits
private func wouldExceedLimits(_ piece: ChessPiece) -> Bool {
let count = piece.color == .white ? whitePieceCount[piece.type] ?? 0 : blackPieceCount[piece.type] ?? 0
switch piece.type {
case .king: return count >= 1
case .queen: return count >= 1
case .rook, .bishop, .knight: return count >= 2
case .pawn: return count >= 8
}
}
/// Recognize a chess piece from a square image
/// - Parameters:
/// - image: CGImage of the chess square
/// - completion: Callback with result (ChessPiece if recognized, nil if empty)
/// - Throws: PieceRecognitionError
func recognizePiece(from image: CGImage) async throws -> ChessPiece? {
// Validate image dimensions // Validate image dimensions
guard image.width > 0, image.height > 0, guard image.width > 0, image.height > 0,
abs(1 - Float(image.width) / Float(image.height)) < 0.1 else { abs(1 - Float(image.width) / Float(image.height)) < 0.1 else {
@ -161,7 +64,6 @@ final class PieceRecognizer {
} }
let handler = VNImageRequestHandler(cgImage: image) let handler = VNImageRequestHandler(cgImage: image)
var classificationResults: [VNClassificationObservation]?
var classificationError: Error? var classificationError: Error?
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
@ -171,7 +73,7 @@ final class PieceRecognizer {
continuation.resume(throwing: error) continuation.resume(throwing: error)
return return
} }
classificationResults = request.results as? [VNClassificationObservation] self.classificationResults = request.results as? [VNClassificationObservation]
continuation.resume() continuation.resume()
} }
request.imageCropAndScaleOption = .centerCrop request.imageCropAndScaleOption = .centerCrop
@ -188,160 +90,33 @@ final class PieceRecognizer {
throw PieceRecognitionError.recognitionFailed(error.localizedDescription) throw PieceRecognitionError.recognitionFailed(error.localizedDescription)
} }
guard let results = classificationResults, guard let results = self.classificationResults,
let topResult = results.first else { let topResult = results.first else {
print("ERROR: No classification results") print("ERROR: No classification results")
throw PieceRecognitionError.recognitionFailed("No results") throw PieceRecognitionError.recognitionFailed("No results")
} }
// Print all results to help diagnose recognition issues // Print results
if let pos = currentPosition { print("\nClassification results:")
print("\nClassification results for \(String(describing: pos)):")
} else {
print("\nClassification results:")
}
for result in results.prefix(3) { for result in results.prefix(3) {
print("- \(result.identifier): \(result.confidence)") print("- \(result.identifier): \(result.confidence)")
} }
// Check for empty squares // Check confidence threshold
if topResult.identifier == "empty_dark" || topResult.identifier == "empty_light" { if topResult.confidence < confidenceThreshold {
if topResult.confidence > 0.9 { print("Low confidence (\(topResult.confidence)) for \(topResult.identifier)")
if let pos = currentPosition { return .empty()
print("\nEmpty square confirmed at \(String(describing: pos))")
}
return nil
}
} }
// Get confidence ratio between top predictions // Try to create a SquareClassification from the label
let secondBestConfidence = results.count > 1 ? results[1].confidence : 0 if let classification = SquareClassification(label: topResult.identifier) {
let confidenceRatio = topResult.confidence / (secondBestConfidence + Float.ulpOfOne) print("Classified as \(topResult.identifier) with confidence \(topResult.confidence)")
return classification
// Adjust confidence based on position-specific knowledge
let adjustedConfidence = adjustConfidence(topResult.confidence,
for: topResult.identifier,
at: currentPosition)
// Much stricter piece recognition:
// 1. Must have very high confidence (>0.98)
// 2. Must have strong separation from second best (>5.0 ratio)
// 3. Must make sense for position
// 4. Must not exceed piece limits
if adjustedConfidence > 0.98 && confidenceRatio > 5.0 {
if let piece = try? createPiece(from: topResult.identifier) {
if isValidPieceForPosition(piece, at: currentPosition) && !wouldExceedLimits(piece) {
updatePieceCount(piece: piece)
return piece
}
}
} }
// Log classification details // Return empty square if classification fails
if let pos = currentPosition { print("Classification failed for \(topResult.identifier) (\(topResult.confidence))")
print("\nClassification rejected at \(String(describing: pos)):") return .empty()
}
print("Top result: \(topResult.identifier) (\(topResult.confidence))")
if results.count > 1 {
print("Second best: \(results[1].identifier) (\(results[1].confidence))")
print("Confidence ratio: \(confidenceRatio)")
}
print("Adjusted confidence: \(adjustedConfidence)")
return nil // Default to empty for unclear cases
}
/// Position where this piece is being recognized
var currentPosition: BoardPosition?
/// Adjust confidence based on position-specific knowledge
private func adjustConfidence(_ confidence: Float, for identifier: String, at position: BoardPosition?) -> Float {
guard let position = currentPosition else { return confidence }
// Parse the piece info
let components = identifier.split(separator: "_")
guard components.count == 2,
let color = PieceColor(rawValue: String(components[0])),
let type = PieceType(rawValue: String(components[1])) else {
return confidence
}
var adjustment: Float = 0.0
// Back rank pieces are more likely to be correct
if position.isBackRank(for: color) {
// Corners should be rooks
if type == .rook && position.isEdgeFile {
adjustment += 0.1
}
// Next to corners should be knights
if type == .knight && (position.file == 1 || position.file == 6) {
adjustment += 0.1
}
// Next to knights should be bishops
if type == .bishop && (position.file == 2 || position.file == 5) {
adjustment += 0.1
}
// Center should be king/queen
if (type == .king || type == .queen) && position.isCenterFile {
adjustment += 0.1
}
}
// Pawns are more likely on their starting ranks
if type == .pawn {
if (color == .white && position.rank == 1) ||
(color == .black && position.rank == 6) {
adjustment += 0.1
}
}
return min(1.0, confidence + adjustment)
}
/// Validate if a piece makes sense for its position
private func isValidPieceForPosition(_ piece: ChessPiece, at position: BoardPosition?) -> Bool {
guard let pos = position else { return true }
// Basic position validation
switch piece.type {
case .king:
// Kings can't be on the first or last rank of opponent's side
if piece.color == .white && pos.rank == 7 { return false }
if piece.color == .black && pos.rank == 0 { return false }
case .pawn:
// Pawns can't be on first or last rank
if pos.rank == 0 || pos.rank == 7 { return false }
// White pawns can't be behind their starting rank
if piece.color == .white && pos.rank > 6 { return false }
// Black pawns can't be behind their starting rank
if piece.color == .black && pos.rank < 1 { return false }
default:
// Other pieces can move freely
break
}
return true
}
/// Helper to create a chess piece from a classification label
private func createPiece(from identifier: String) throws -> ChessPiece {
// Validate it's not an empty square
guard !identifier.starts(with: "empty_") else {
print("ERROR: Cannot create piece from empty square label: \(identifier)")
throw PieceRecognitionError.invalidInput
}
let components = identifier.split(separator: "_")
guard components.count == 2,
let color = PieceColor(rawValue: String(components[0])),
let type = PieceType(rawValue: String(components[1])) else {
print("ERROR: Invalid piece label format: \(identifier)")
throw PieceRecognitionError.invalidInput
}
return ChessPiece(type: type, color: color)
} }
/// Preprocess an image for recognition /// Preprocess an image for recognition
@ -351,15 +126,28 @@ final class PieceRecognizer {
let ciImage = CIImage(cgImage: image) let ciImage = CIImage(cgImage: image)
// Apply preprocessing filters // Apply preprocessing filters
let processed = ciImage // First pass: Enhance contrast and edges
var processed = ciImage
.applyingFilter("CIColorControls", parameters: [ .applyingFilter("CIColorControls", parameters: [
kCIInputContrastKey: 1.1, "inputContrast": 1.3,
kCIInputBrightnessKey: 0.0, "inputBrightness": 0.0,
kCIInputSaturationKey: 1.1 "inputSaturation": 1.0
]) ])
.applyingFilter("CIUnsharpMask", parameters: [ .applyingFilter("CIUnsharpMask", parameters: [
kCIInputRadiusKey: 1.0, "inputRadius": 2.0,
kCIInputIntensityKey: 0.5 "inputIntensity": 0.8
])
// Second pass: Reduce noise and enhance details
processed = processed
.applyingFilter("CINoiseReduction", parameters: [
"inputNoiseLevel": 0.2,
"inputSharpness": 0.6
])
.applyingFilter("CIColorControls", parameters: [
"inputContrast": 1.2,
"inputBrightness": 0.0,
"inputSaturation": 1.0
]) ])
// Convert back to CGImage // Convert back to CGImage
@ -370,28 +158,3 @@ final class PieceRecognizer {
return outputImage return outputImage
} }
} }
// MARK: - BoardPosition Extensions
extension BoardPosition {
/// Initialize from file and rank indices
init(file: Int, rank: Int) {
self.file = file
self.rank = rank
}
/// Whether this position is on the edge of the board (files a or h)
var isEdgeFile: Bool {
return file == 0 || file == 7
}
/// Whether this position is on the back rank for the given color
func isBackRank(for color: PieceColor) -> Bool {
return (color == .white && rank == 0) || (color == .black && rank == 7)
}
/// Whether this position is in the center files (d or e)
var isCenterFile: Bool {
return file == 3 || file == 4
}
}

File diff suppressed because it is too large Load diff

View file

@ -1,29 +1,34 @@
# Current Task # Active Context
Working on chess piece recognition with the updated ML model that includes empty square detection.
# Recent Changes ## Current Task
1. Empty Square Detection: - Fixed model integration issues
- Handle empty_dark/empty_light classes - Updated code to match model categories
- Exact class name matching - Simplified classification system
- Proper error handling
2. Position Validation: ## Recent Changes
- Allow moved pieces 1. SquareClassification.swift:
- Essential rules only - Exact category mapping:
- Piece count tracking ```swift
white_pawn, white_knight, white_bishop, white_rook, white_queen, white_king,
black_pawn, black_knight, black_bishop, black_rook, black_queen, black_king
```
- Default to empty square for unrecognized labels
- Removed background variations from labels
3. Error Handling: 2. PieceRecognizer.swift:
- Better error messages - Using VNCoreMLModel consistently
- Clear logging - Simplified error handling
- Fixed optional unwrapping - Removed unnecessary piece counting
- Using Vision framework for classification
# Next Steps ## Next Steps
1. Recognition Tuning: 1. Test model integration:
- Fine-tune empty square detection - Verify model loads correctly
- Adjust confidence thresholds - Check classification accuracy
- Improve position validation - Monitor confidence levels
2. Model Training: ## Current Issues
- Add more empty square examples Fixed:
- Include different board styles - Model pipeline error
- Improve piece variety - Category mismatches
- Classification handling

View file

@ -9,12 +9,4 @@ Current position:
1 wR wN .. .. wK .. wN wR 1 wR wN .. .. wK .. wN wR
a b c d e f g h a b c d e f g h
Key differences from recognition:
4. Some pieces missing from recognition
5. Some pieces misidentified
Recognition issues to fix:
1. Need to handle moved pieces (not just starting position)
2. Better empty square detection
3. Improve confidence thresholds
4. Validate complete position

View file

@ -1,60 +1,89 @@
# Product Overview # Product Context
ChessPrism is a macOS application that captures and analyzes chess positions from the screen in real-time.
## Project Purpose
ChessPrism is a macOS application designed to:
1. Capture chess.com game windows
2. Detect and analyze chess positions in real-time
3. Generate FEN strings for position analysis
4. Track moves and game progress
## Core Features ## Core Features
1. Board Detection 1. Screen Capture
- Automatic chessboard location - Automatic chess.com window detection
- Perspective and size handling - Real-time board monitoring
- Multi-board support planned - Configurable capture settings
2. Piece Recognition 2. Board Analysis
- ML-based piece classification - Accurate piece detection
- Empty square detection - 12 piece categories:
- Position-aware confidence adjustments * 6 white pieces (pawn to king)
* 6 black pieces (pawn to king)
3. Position Analysis
- FEN string generation
- Position validation - Position validation
- Move tracking (planned) - FEN string generation
## Current Challenges 3. Machine Learning
- Vision-based Core ML model
- Direct category mapping
- High confidence threshold (0.75)
- Fast inference time
### Recognition Features ## User Experience Goals
1. Empty Square Detection 1. Reliability
- Explicit empty_dark/empty_light classes - Accurate piece detection
- Direct square color recognition - Consistent board recognition
- High confidence classification - Robust error handling
2. Piece Recognition 2. Performance
- Accurate piece type detection - Real-time analysis
- Color differentiation - Low resource usage
- Position-aware confidence - Smooth capture
3. Performance Optimization 3. Usability
- Fast-path empty detection - Automatic window detection
- Efficient classification flow - Minimal setup required
- Resource-aware processing - Clear feedback
## Future Improvements ## Current Status
1. Working Features
- Screen capture system
- Board detection
- Piece recognition
- FEN generation
### Short Term 2. Recent Improvements
1. Recognition Enhancement - Simplified classification system
- Fine-tune confidence thresholds - Direct category mapping
- Validate square colors - Vision framework integration
- Improve error messages - Improved error handling
2. Position Analysis 3. Known Limitations
- Move validation - Requires chess.com's default board theme
- Game state tracking - macOS 12.3+ requirement
- Historical context - Screen capture permissions needed
### Long Term ## Future Enhancements
1. Advanced Features 1. Short Term
- Move detection - Monitor classification accuracy
- Game recording - Fine-tune confidence threshold
- Multiple board styles - Improve error reporting
2. User Experience 2. Long Term
- Confidence visualization - Support for multiple board themes
- Manual corrections - Game analysis integration
- Custom training - Move suggestion system
## Technical Requirements
1. System
- macOS 12.3 or later
- Metal-capable GPU
- Screen recording permissions
2. Dependencies
- Vision framework
- Core ML
- ScreenCaptureKit
3. Performance Targets
- 30 FPS capture
- Sub-second analysis
- Low CPU/GPU usage

139
cline_docs/sample.py Normal file
View file

@ -0,0 +1,139 @@
import os
import requests
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
from pathlib import Path
from io import BytesIO
import itertools
class ChessAssetProcessor:
def __init__(self):
self.base_dir = Path.cwd()
self.raw_dir = self.base_dir / 'raw'
self.training_dir = self.base_dir / 'training'
# URLs for pieces
self.piece_base_url = 'https://www.chess.com/chess-themes/pieces/neo/300'
# Piece configurations - matches the Swift enum exactly
self.pieces = {
'white': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king'],
'black': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king']
}
# Square colors - used for background variations only
self.square_colors = {
'light': '#eeeed2',
'dark': '#759656',
'light_highlighted': '#f6f68d',
'dark_highlighted': '#bdcc49'
}
# Border configurations
self.border_configs = [
{}, # No borders
{'top': True},
{'bottom': True},
{'left': True},
{'top': True, 'left': True},
{'bottom': True, 'left': True},
]
self._setup_directories()
def _setup_directories(self):
"""Create directory structure for training data"""
(self.raw_dir / 'pieces').mkdir(parents=True, exist_ok=True)
self.training_dir.mkdir(exist_ok=True)
# Create directories for each piece type
for color in ['white', 'black']:
for piece in self.pieces[color]:
(self.training_dir / f"{color}_{piece}").mkdir(parents=True, exist_ok=True)
def hex_to_rgb(self, hex_color):
"""Convert hex color to RGB tuple"""
hex_color = hex_color.lstrip('#')
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
def create_base_square(self, color, size=100):
"""Create a square with specified color"""
rgb_color = self.hex_to_rgb(color)
return Image.new('RGB', (size, size), rgb_color)
def add_borders(self, image, borders, border_color=(0, 0, 0)):
"""Add borders according to configuration"""
w, h = image.size
result = image.copy()
draw = ImageDraw.Draw(result)
border_size = 5
if borders.get('top'):
draw.line([(0, 0), (w-1, 0)], fill=border_color, width=border_size)
if borders.get('bottom'):
draw.line([(0, h-1), (w-1, h-1)], fill=border_color, width=border_size)
if borders.get('left'):
draw.line([(0, 0), (0, h-1)], fill=border_color, width=border_size)
return result
def download_piece(self, color, piece, target_size):
"""Download a specific chess piece"""
piece_letter = piece[0] if piece != 'knight' else 'n'
piece_url = f"{self.piece_base_url}/{color[0]}{piece_letter}.png"
response = requests.get(piece_url)
if response.status_code == 200:
piece_img = Image.open(BytesIO(response.content)).convert('RGBA')
return piece_img.resize((target_size, target_size), Image.Resampling.LANCZOS)
return None
def create_training_data(self):
"""Create comprehensive training dataset"""
square_size = 100
print("Downloading pieces and creating variations...")
for color in ['white', 'black']:
for piece in self.pieces[color]:
piece_dir = self.training_dir / f"{color}_{piece}"
print(f"\nProcessing {color} {piece}...")
# Download piece
piece_img = self.download_piece(color, piece, square_size)
if piece_img is None:
print(f"Failed to download {color} {piece}")
continue
# Create variations with different backgrounds
variation_count = 0
for bg_name, bg_color in self.square_colors.items():
# Create base square
base_square = self.create_base_square(bg_color, square_size)
# Add border variations
for border_config in self.border_configs:
# Add borders
bordered = self.add_borders(base_square, border_config)
# Convert to RGBA for composition
bordered_rgba = bordered.convert('RGBA')
# Combine with piece
combined = Image.alpha_composite(bordered_rgba, piece_img)
# Generate filename
border_desc = '_'.join(k for k,v in border_config.items() if v)
filename = f"{bg_name}_{border_desc}_{variation_count}.png" if border_desc else f"{bg_name}_{variation_count}.png"
# Save image
combined.save(piece_dir / filename)
variation_count += 1
print(f"Created {variation_count} variations for {color} {piece}")
def main():
processor = ChessAssetProcessor()
processor.create_training_data()
if __name__ == '__main__':
main()

View file

@ -1,81 +1,93 @@
# System Architecture # System Patterns
## Piece Recognition Pipeline ## Model Architecture
1. Core ML Integration
- Vision framework for image handling
- Direct category mapping
- No intermediate transformations
2. Classification Flow
```
Image → VNImageRequestHandler → VNCoreMLRequest → VNClassificationObservation → SquareClassification
```
3. Category System
- 12 piece categories:
* white_pawn to white_king
* black_pawn to black_king
- Empty square fallback
- No background variations in model
## Processing Pipeline
1. Board Detection 1. Board Detection
- VNDetectRectanglesRequest for board location - VNDetectRectanglesRequest
- Aspect ratio and size validation - Aspect ratio validation
- Square extraction with equal dimensions - Square extraction
2. Image Preprocessing 2. Image Processing
- Contrast and brightness adjustment - Contrast enhancement
- Unsharp mask for edge enhancement - Edge sharpening
- Consistent image scaling - Noise reduction
- Center crop
3. ML Classification 3. Classification
- CoreML model prediction - Vision framework integration
- Confidence score analysis - Confidence threshold
- Position-based adjustments - Error handling
## Recognition Patterns ## Code Organization
1. Recognition Layer
```
PieceRecognizer
├── Model loading
├── Image preprocessing
└── Classification handling
```
### Square Classification 2. Model Layer
1. Empty Square Detection ```
- Exact empty_dark/empty_light matching SquareClassification
- High confidence threshold (>0.9) ├── Category mapping
- Early detection and return ├── Piece type/color
└── Empty square handling
```
2. Piece Recognition 3. Core Components
- Strict label format validation ```
- Position-based confidence adjustment BoardDetector
- Piece count tracking ├── Rectangle detection
├── Square extraction
└── Position validation
```
3. Error Prevention ## Data Flow
- Empty square validation 1. Capture
- Label format checking ```
- Position rule enforcement ScreenCapture → Raw Image → Board Rectangle
```
### Classification Flow 2. Processing
1. Input Validation ```
- Image dimensions check Board Rectangle → Individual Squares → Preprocessed Images
- Model availability check ```
- Configuration setup
2. Square Analysis 3. Classification
- Empty square check first ```
- Piece classification second Preprocessed Images → ML Model → Piece Categories → Chess Position
- Position validation last ```
3. Confidence Checks ## Key Patterns
- Empty squares: >0.9 1. Direct Integration
- Pieces: >0.98 with >5.0 ratio - Vision framework throughout
- Position adjustments - No intermediate conversions
- Consistent image handling
4. Error Handling 2. Error Handling
- Clear error messages - Early validation
- Graceful fallbacks
- Detailed logging - Detailed logging
- Safe fallbacks
## Validation Patterns 3. Performance
1. Piece Count Rules - Shared CIContext
- Maximum 1: king, queen - Efficient image processing
- Maximum 2: rooks, bishops, knights - Optimized model loading
- Maximum 8: pawns
- Track by color and type
2. Position Rules
- Kings: not on opponent's back rank
- Pawns: no backward movement
- All pieces: within board bounds
- All pieces: valid movement patterns
3. Piece Tracking
- Maximum piece counts
- Color-specific tracking
- Total position validation
- Captured piece limits
4. Recognition Flow
- Empty square detection first
- Piece classification second
- Position validation last
- Clear error reporting

View file

@ -1,69 +1,88 @@
# Technologies Used # Technical Context
## Core ML & Vision
- ChessPieceClassifier.mlmodel for piece recognition
- VNCoreMLModel for image classification
- Vision framework for board detection
## Image Processing
- CoreImage for preprocessing
- CIColorControls and CIUnsharpMask filters
- CGImage for image manipulation
## Development Environment ## Development Environment
- Xcode for Swift development - macOS Application
- Create ML for model training - Swift & SwiftUI
- SwiftUI for UI components - Xcode 14+
- Target: macOS 12.3+
# Technical Constraints ## Core Technologies
1. Vision Framework
- VNDetectRectanglesRequest for board detection
- VNCoreMLRequest for piece classification
- VNImageRequestHandler for image processing
## ML Model Capabilities 2. Core ML Model
1. Classification Types: - Name: ChessPieceClassifier.mlmodel
- Pieces: pawn, rook, knight, bishop, queen, king - Input: RGB/RGBA images
- Colors: black, white - Output: Classification label
- Empty squares: dark, light - Categories (exact names):
- Label formats: color_piece, empty_color ```swift
white_pawn, white_knight, white_bishop, white_rook, white_queen, white_king,
black_pawn, black_knight, black_bishop, black_rook, black_queen, black_king
```
2. Recognition Features: 3. ScreenCaptureKit
- Multi-class classification - Window capture at 30 FPS
- Per-class confidence scores - Configurable cursor visibility
- Position-aware validation - Chess.com window detection
- Piece count tracking
## Processing Requirements ## Image Processing
1. Image Requirements: 1. Preprocessing Pipeline
- Square dimensions (1:1 ±10%) - Contrast enhancement (1.3x)
- Non-zero dimensions - Edge sharpening
- Center-cropped squares - Noise reduction
- Clear piece visibility - Color normalization
2. Recognition Rules: 2. Square Extraction
- Empty squares: exact class match with >0.9 confidence - Aspect ratio validation
- Pieces: strict format with >0.98 confidence - Size normalization
- Separation ratio: >5.0 between predictions - Center crop
- Position validation: essential rules only
3. Error Prevention: ## Model Integration
- Early empty square detection 1. Loading
- Strict label validation ```swift
- Safe optional handling let config = MLModelConfiguration()
- Clear error messages config.computeUnits = .all
let model = try MLModel(contentsOf: modelURL)
let vnModel = try VNCoreMLModel(for: model)
```
## Performance Considerations 2. Classification
1. Processing Flow: ```swift
- Early validation checks let request = VNCoreMLRequest(model: vnModel)
- Fast empty square detection request.imageCropAndScaleOption = .centerCrop
- Efficient error handling ```
- Quick rejection paths
2. Resource Optimization: 3. Result Handling
- GPU acceleration for ML - Confidence threshold: 0.75
- Minimal preprocessing - Empty square fallback
- Optimized validation - Direct category mapping
- Efficient logging
# Development Setup ## Dependencies
1. Clone repository - Foundation
2. Open ChessPrism.xcodeproj - Vision
3. Build and run on macOS - CoreML
4. Model at ChessPrism/ChessPieceClassifier.mlmodel - CoreImage
- ScreenCaptureKit
- SwiftUI
## Error Handling
- Invalid dimensions
- Model loading failures
- Recognition errors
- Low confidence results
## File Organization
```
ChessPrism/
├── Models/
│ ├── SquareClassification.swift # Model output mapping
│ └── ChessPosition.swift # Board state
├── Recognition/
│ ├── PieceRecognizer.swift # ML integration
│ ├── FenGenerator.swift # Position encoding
│ └── MoveDetector.swift # Move analysis
└── Core/
├── BoardDetector.swift # Square extraction
└── ScreenCapture.swift # Window capture