Spaces:
Sleeping
Sleeping
File size: 2,932 Bytes
e0152b8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import shutil
from smolagents import Tool
import requests
class ChessEngineLocatorTool(Tool):
name = "locate_chess_engine"
description = "Locates the installed Stockfish chess engine binary on the system."
inputs = {}
output_type = "string"
def forward(self) -> str:
engine_path = shutil.which("stockfish")
return engine_path or "Stockfish engine not found."
class ChessboardImageToFENTool(Tool):
name = "image_to_fen"
description = "Analyzes a chessboard image and returns its FEN representation."
inputs = {
"image_url": {
"type": "string",
"description": "A public image URL or a base64 data URL containing a chessboard photo."
}
}
output_type = "string"
def __init__(self, endpoint_url: str, hf_token: str, **kwargs):
super().__init__(**kwargs)
self.endpoint_url = endpoint_url
self.hf_token = hf_token
def forward(self, image_url: str) -> str:
headers = {
"Authorization": f"Bearer {self.hf_token}",
"Content-Type": "application/json"
}
payload = {
"inputs": [
{
"role": "user",
"content": [
{"type": "input_text", "text": "Extract chess piece positions from this image."},
{"type": "input_image", "image_url": image_url}
]
},
{
"role": "user",
"content": [
{
"type": "input_text",
"text": (
"Using the extracted pieces, list each one with its position in format <piece><square> (e.g., Kd4).\n"
"Use uppercase for white, lowercase for black. Output only the lines with positions."
)
}
]
}
]
}
response = requests.post(self.endpoint_url, headers=headers, json=payload)
response.raise_for_status()
raw_text = response.json().get("generated_text", "")
squares = {}
for line in raw_text.splitlines():
line = line.strip()
if len(line) == 3:
squares[line[1:3]] = line[0]
fen_rows = []
for rank in range(8, 0, -1):
row = ""
empty = 0
for file in "abcdefgh":
pos = f"{file}{rank}"
if pos in squares:
if empty > 0:
row += str(empty)
empty = 0
row += squares[pos]
else:
empty += 1
if empty:
row += str(empty)
fen_rows.append(row)
return "/".join(fen_rows) |