app.py

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