Test_Magus / prechess.py
SergeyO7's picture
Create prechess.py
f706229 verified
@tool
def PNG2FENTool(png_file: str) -> str:
"""Tool for converting a PNG file containing a chess board to a FEN position string.
Args:
png_file (str): The path to the PNG file.
Returns:
str: The FEN position string representing the chess board.
"""
# Raises:
# - FileNotFoundError:
# If the PNG file does not exist.
# - ValueError:
# If the PNG file cannot be processed or does not contain a valid chess board.
try:
# Open and preprocess image with modern Pillow
img = Image.open(png_file)
img = ImageOps.exif_transpose(img).convert("L")
# Use LANCZOS instead of ANTIALIAS
img = img.resize((img.width*2, img.height*2), Image.Resampling.LANCZOS)
# Save temp file for OCR
temp_path = "chess_temp.png"
img.save(temp_path)
# Perform OCR
import easyocr
reader = easyocr.Reader(['en'])
result = reader.readtext(png_file, detail=0)
fen_candidates = [text for text in result if validate_fen_format(text)]
if not fen_candidates:
raise ValueError("No valid FEN found in image")
return fen_candidates[0]
except Exception as e:
raise ValueError(f"OCR processing failed: {str(e)}")
# try:
# # Open the PNG file using PIL
# image = Image.open(png_file)
#
# # Use pytesseract to extract text from the image
# text = pytesseract.image_to_string(image)
#
# # Process the extracted text to get the FEN position string
# fen_position = process_text_to_fen(text)
#
# return fen_position
#
except FileNotFoundError:
raise FileNotFoundError("PNG file not found.")
#
# except Exception as e:
# raise ValueError("Error processing PNG file: " + str(e))
def process_text_to_fen(text):
"""
Processes the extracted text from the image to obtain the FEN position string.
Parameters:
- text: str
The extracted text from the image.
Returns:
- str:
The FEN position string representing the chess board.
Raises:
- ValueError:
If the extracted text does not contain a valid chess board.
"""
# Process the text to remove any unnecessary characters or spaces
processed_text = text.strip().replace("\n", "").replace(" ", "")
# Check if the processed text matches the expected format of a FEN position string
if not validate_fen_format(processed_text):
raise ValueError("Invalid chess board.")
return processed_text
def validate_fen_format(fen_string):
"""
Validates if a given string matches the format of a FEN (Forsyth–Edwards Notation) position string.
Parameters:
- fen_string: str
The string to be validated.
Returns:
- bool:
True if the string matches the FEN format, False otherwise.
"""
# FEN format: 8 sections separated by '/'
sections = fen_string.split("/")
if len(sections) != 8:
return False
# Check if each section contains valid characters
for section in sections:
if not validate_section(section):
return False
return True
def validate_section(section):
"""
Validates if a given section of a FEN (Forsyth–Edwards Notation) position string contains valid characters.
Parameters:
- section: str
The section to be validated.
Returns:
- bool:
True if the section contains valid characters, False otherwise.
"""
# Valid characters: digits 1-8 or letters 'r', 'n', 'b', 'q', 'k', 'p', 'R', 'N', 'B', 'Q', 'K', 'P'
valid_chars = set("12345678rnbqkpRNBQKP")
return all(char in valid_chars for char in section)
import chess
import chess.engine
class ChessEngineTool(Tool):
name = "chess_engine"
description = "Analyzes a chess position (FEN) with Stockfish and returns the best move."
inputs = {
"fen": {"type": "string", "description": "FEN string of the position."},
"time_limit": {"type": "number", "description": "Time in seconds for engine analysis.", "nullable": True}
}
output_type = "string"
def forward(self, fen: str, time_limit: float = 0.1) -> str:
# figure out where the binary actually is
sf_bin = shutil.which("stockfish") or "/usr/games/stockfish"
if not sf_bin:
raise RuntimeError(
f"Cannot find stockfish on PATH or at /usr/games/stockfish. "
"Did you install it in apt.txt or via apt-get?"
)
board = chess.Board(fen)
engine = chess.engine.SimpleEngine.popen_uci(sf_bin)
result = engine.play(board, chess.engine.Limit(time=time_limit))
engine.quit()
return board.san(result.move)