ChessPrism/ChessPrism.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings

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

ChessPrism/ChessPrism/BoardDetector.swift

@@ -95,11 +95,11 @@ final class BoardDetector {
             throw BoardDetectionError.invalidBoardDimensions
         }
         
-        var position = ChessPosition()
-        var recognizedPieces = 0
-        
         // Reset piece counts before starting new scan
-        pieceRecognizer.resetPieceCounts()
+        pieceRecognizer.resetCounts()
+        
+        var squares = [BoardPosition: SquareClassification]()
+        var pieceCount = 0
         
         // Process each square
         for rank in 0..<8 {
@@ -107,33 +107,27 @@ final class BoardDetector {
                 let boardPosition = BoardPosition(file: file, rank: rank)
                 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)
+        guard let position = ChessPosition(fen: fen) else {
+            print("ERROR: Failed to create position from FEN")
+            throw BoardDetectionError.imageProcessingFailed
+        }
         
-        print("\nPieces recognized: \(recognizedPieces)")
+        print("\nPieces recognized: \(pieceCount)")
-        print("\nRecognized position:")
+        print("\nGenerated FEN: \(fen)")
-        for rank in (0...7).reversed() {
-            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")
         
         guard position.isValid else {
             print("\nERROR: Invalid chess position")

ChessPrism/ChessPrism/ChessPieceClassifier.mlmodelc/metadata.json

@@ -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"
-  }
-]

ChessPrism/ChessPrism/ChessPrism.entitlements

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

ChessPrism/ChessPrism/Models/ChessPosition.swift

@@ -41,10 +41,25 @@ struct ChessPiece: Equatable {
 }
 
 /// Represents a position on the chess board
-struct BoardPosition: Equatable {
+struct BoardPosition: Equatable, Hashable {
     let file: Int // 0-7 for a-h
     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")
     init?(algebraic: String) {
         guard algebraic.count == 2,

ChessPrism/ChessPrism/Models/SquareClassification.swift

@@ -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)
+        }
+    }
+}

ChessPrism/ChessPrism/Recognition/FenGenerator.swift

@@ -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()
+    }
+}

ChessPrism/ChessPrism/Recognition/MoveDetector.swift

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

ChessPrism/ChessPrism/Recognition/PieceRecognizer.swift

@@ -10,7 +10,6 @@ enum PieceRecognitionError: Error {
     case recognitionFailed(String)
     case lowConfidence
     case invalidInput
-    case invalidPosition(String)
 }
 
 /// A class responsible for recognizing chess pieces from images
@@ -23,9 +22,11 @@ final class PieceRecognizer {
     /// Shared CIContext for image processing
     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
     
@@ -34,7 +35,7 @@ final class PieceRecognizer {
         
         let bundle = Bundle.main
         
-        // Load model from bundle
+        // Load model
         guard let modelURL = bundle.url(forResource: "ChessPieceClassifier", withExtension: "mlmodelc") else {
             print("ERROR: Model not found in bundle at \(bundle.bundlePath)")
             throw PieceRecognitionError.modelLoadError
@@ -50,109 +51,11 @@ final class PieceRecognizer {
             print("ERROR: Failed to load model - \(error)")
             throw PieceRecognitionError.modelLoadError
         }
-        
-        resetPieceCounts()
     }
     
     // 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
         guard image.width > 0, image.height > 0,
               abs(1 - Float(image.width) / Float(image.height)) < 0.1 else {
@@ -161,7 +64,6 @@ final class PieceRecognizer {
         }
         
         let handler = VNImageRequestHandler(cgImage: image)
-        var classificationResults: [VNClassificationObservation]?
         var classificationError: Error?
         
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
@@ -171,7 +73,7 @@ final class PieceRecognizer {
                     continuation.resume(throwing: error)
                     return
                 }
-                classificationResults = request.results as? [VNClassificationObservation]
+                self.classificationResults = request.results as? [VNClassificationObservation]
                 continuation.resume()
             }
             request.imageCropAndScaleOption = .centerCrop
@@ -188,160 +90,33 @@ final class PieceRecognizer {
             throw PieceRecognitionError.recognitionFailed(error.localizedDescription)
         }
         
-        guard let results = classificationResults,
+        guard let results = self.classificationResults,
               let topResult = results.first else {
             print("ERROR: No classification results")
             throw PieceRecognitionError.recognitionFailed("No results")
         }
         
-        // Print all results to help diagnose recognition issues
+        // Print results
-        if let pos = currentPosition {
-            print("\nClassification results for \(String(describing: pos)):")
-        } else {
         print("\nClassification results:")
-        }
         for result in results.prefix(3) {
             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
@@ -351,15 +126,28 @@ final class PieceRecognizer {
         let ciImage = CIImage(cgImage: image)
         
         // Apply preprocessing filters
-        let processed = ciImage
+        // First pass: Enhance contrast and edges
+        var processed = ciImage
             .applyingFilter("CIColorControls", parameters: [
-                kCIInputContrastKey: 1.1,
+                "inputContrast": 1.3,
-                kCIInputBrightnessKey: 0.0,
+                "inputBrightness": 0.0,
-                kCIInputSaturationKey: 1.1
+                "inputSaturation": 1.0
             ])
             .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
@@ -370,28 +158,3 @@ final class PieceRecognizer {
         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
-    }
-}

cline_docs/activeContext.md

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

cline_docs/chessboard.txt

@@ -9,12 +9,4 @@ Current position:
 1 wR wN .. .. wK .. wN wR 
   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

cline_docs/productContext.md

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

cline_docs/sample.py

@@ -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()

cline_docs/systemPatterns.md

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

cline_docs/techContext.md

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