from ultralytics import YOLO import cv2 from stockfish import Stockfish import os import numpy as np import streamlit as st import requests # Constants FEN_MAPPING = { "black-pawn": "p", "black-rook": "r", "black-knight": "n", "black-bishop": "b", "black-queen": "q", "black-king": "k", "white-pawn": "P", "white-rook": "R", "white-knight": "N", "white-bishop": "B", "white-queen": "Q", "white-king": "K" } GRID_BORDER = 10 # Border size in pixels GRID_SIZE = 204 # Effective grid size (10px to 214px) BLOCK_SIZE = GRID_SIZE // 8 # Each block is ~25px X_LABELS = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'] # Labels for x-axis (a to h) Y_LABELS = [8, 7, 6, 5, 4, 3, 2, 1] # Reversed labels for y-axis (8 to 1) # Functions def get_grid_coordinate(pixel_x, pixel_y): """ Function to determine the grid coordinate of a pixel, considering a 10px border and the grid where bottom-left is (a, 1) and top-left is (h, 8). """ # Grid settings border = 10 # 10px border grid_size = 204 # Effective grid size (10px to 214px) block_size = grid_size // 8 # Each block is ~25px x_labels = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'] # Labels for x-axis (a to h) y_labels = [8, 7, 6, 5, 4, 3, 2, 1] # Reversed labels for y-axis (8 to 1) # Adjust pixel_x and pixel_y by subtracting the border (grid starts at pixel 10) adjusted_x = pixel_x - border adjusted_y = pixel_y - border # Check bounds if adjusted_x < 0 or adjusted_y < 0 or adjusted_x >= grid_size or adjusted_y >= grid_size: return "Pixel outside grid bounds" # Determine the grid column and row x_index = adjusted_x // block_size y_index = adjusted_y // block_size if x_index < 0 or x_index >= len(x_labels) or y_index < 0 or y_index >= len(y_labels): return "Pixel outside grid bounds" # Convert indices to grid coordinates x_index = adjusted_x // block_size # Determine the column index (0-7) y_index = adjusted_y // block_size # Determine the row index (0-7) # Convert row index to the correct label, with '8' at the bottom y_labeld = y_labels[y_index] # Correct index directly maps to '8' to '1' x_label = x_labels[x_index] y_label = 8 - y_labeld + 1 return f"{x_label}{y_label}" def predict_next_move(fen, stockfish): """ Predict the next move using Stockfish. """ if stockfish.is_fen_valid(fen): stockfish.set_fen_position(fen) else: return "Invalid FEN notation!" best_move = stockfish.get_best_move() ans = transform_string(best_move) return f"The predicted next move is: {ans}" if best_move else "No valid move found (checkmate/stalemate)." # def download_stockfish(): # url = "https://drive.google.com/file/d/18pkwBVc13fgKP3LzrTHE4yzhjyGJexlR/view?usp=sharing" # Replace with the actual link # file_name = "stockfish-windows-x86-64-avx2.exe" # if not os.path.exists(file_name): # print(f"Downloading {file_name}...") # response = requests.get(url, stream=True) # with open(file_name, "wb") as file: # for chunk in response.iter_content(chunk_size=1024): # if chunk: # file.write(chunk) # print(f"{file_name} downloaded successfully.") def process_image(image_path): # Ensure output directory exists if not os.path.exists('output'): os.makedirs('output') # Load the segmentation model segmentation_model = YOLO("segmentation.pt") # Run inference to get segmentation results results = segmentation_model.predict( source=image_path, conf=0.8 # Confidence threshold ) # Initialize variables for the segmented mask and bounding box segmentation_mask = None bbox = None for result in results: if result.boxes.conf[0] >= 0.8: # Filter results by confidence segmentation_mask = result.masks.data.cpu().numpy().astype(np.uint8)[0] bbox = result.boxes.xyxy[0].cpu().numpy() # Get the bounding box coordinates break if segmentation_mask is None: print("No segmentation mask with confidence above 0.8 found.") return None # Load the image image = cv2.imread(image_path) # Resize segmentation mask to match the input image dimensions segmentation_mask_resized = cv2.resize(segmentation_mask, (image.shape[1], image.shape[0])) # Extract bounding box coordinates if bbox is not None: x1, y1, x2, y2 = bbox # Crop the segmented region based on the bounding box cropped_segment = image[int(y1):int(y2), int(x1):int(x2)] # Save the cropped segmented image cropped_image_path = 'output/cropped_segment.jpg' cv2.imwrite(cropped_image_path, cropped_segment) print(f"Cropped segmented image saved to {cropped_image_path}") st.image(cropped_segment, caption="Uploaded Image", use_column_width=True) # Return the cropped image return cropped_segment def transform_string(input_str): # Remove extra spaces and convert to lowercase input_str = input_str.strip().lower() # Check if input is valid if len(input_str) != 4 or not input_str[0].isalpha() or not input_str[1].isdigit() or \ not input_str[2].isalpha() or not input_str[3].isdigit(): return "Invalid input" # Define mappings letter_mapping = { 'a': 'h', 'b': 'g', 'c': 'f', 'd': 'e', 'e': 'd', 'f': 'c', 'g': 'b', 'h': 'a' } number_mapping = { '1': '8', '2': '7', '3': '6', '4': '5', '5': '4', '6': '3', '7': '2', '8': '1' } # Transform string result = "" for i, char in enumerate(input_str): if i % 2 == 0: # Letters result += letter_mapping.get(char, "Invalid") else: # Numbers result += number_mapping.get(char, "Invalid") # Check for invalid transformations if "Invalid" in result: return "Invalid input" return result # Streamlit app def main(): # download_stockfish() st.title("Chessboard Position Detection and Move Prediction") os.chmod("/home/user/app/stockfish-ubuntu-x86-64-sse41-popcnt", 0o755) st.write(os.getcwd()) # User uploads an image or captures it from their camera image_file = st.camera_input("Capture a chessboard image") or st.file_uploader("Upload a chessboard image", type=["jpg", "jpeg", "png"]) if image_file is not None: # Save the image to a temporary file temp_dir = "temp_images" os.makedirs(temp_dir, exist_ok=True) temp_file_path = os.path.join(temp_dir, "uploaded_image.jpg") with open(temp_file_path, "wb") as f: f.write(image_file.getbuffer()) # Process the image using its file path processed_image = process_image(temp_file_path) if processed_image is not None: # Resize the image to 224x224 processed_image = cv2.resize(processed_image, (224, 224)) height, width, _ = processed_image.shape # Initialize the YOLO model model = YOLO("standard.pt") # Replace with your trained model weights file # Run detection results = model.predict(source=processed_image, save=False, save_txt=False, conf=0.6) # Initialize the board for FEN (empty rows represented by "8") board = [["8"] * 8 for _ in range(8)] # Extract predictions and map to FEN board for result in results[0].boxes: x1, y1, x2, y2 = result.xyxy[0].tolist() class_id = int(result.cls[0]) class_name = model.names[class_id] # Convert class_name to FEN notation fen_piece = FEN_MAPPING.get(class_name, None) if not fen_piece: continue # Calculate the center of the bounding box center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 # Convert to integer pixel coordinates pixel_x = int(center_x) pixel_y = int(height - center_y) # Flip Y-axis for generic coordinate system # Get grid coordinate grid_position = get_grid_coordinate(pixel_x, pixel_y) if grid_position != "Pixel outside grid bounds": file = ord(grid_position[0]) - ord('a') # Column index (0-7) rank = int(grid_position[1]) - 1 # Row index (0-7) # Place the piece on the board board[7 - rank][file] = fen_piece # Flip rank index for FEN # Generate the FEN string fen_rows = [] for row in board: fen_row = "" empty_count = 0 for cell in row: if cell == "8": empty_count += 1 else: if empty_count > 0: fen_row += str(empty_count) empty_count = 0 fen_row += cell if empty_count > 0: fen_row += str(empty_count) fen_rows.append(fen_row) position_fen = "/".join(fen_rows) # Ask the user for the next move side move_side = st.selectbox("Select the side to move:", ["w (White)", "b (Black)"]) move_side = "w" if move_side.startswith("w") else "b" # Append the full FEN string continuation fen_notation = f"{position_fen} {move_side} - - 0 0" st.subheader("Generated FEN Notation:") st.code(fen_notation) # Initialize the Stockfish engine stockfish_path = os.path.join(os.getcwd(), "stockfish-ubuntu-x86-64-sse41-popcnt") stockfish = Stockfish( path=stockfish_path, depth=15, parameters={"Threads": 2, "Minimum Thinking Time": 30} ) # Predict the next move next_move = predict_next_move(fen_notation, stockfish) st.subheader("Stockfish Recommended Move:") st.write(next_move) else: st.error("Failed to process the image. Please try again.") if __name__ == "__main__": main()