Spaces:
Running
Running
File size: 4,498 Bytes
3f9559b | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | #!/usr/bin/env python3
"""Faithful reproduction of Preprocess.ipynb crop pipeline on samples/board.png,
then classify the 64 squares and compare to the known FEN."""
import cv2 as cv
import numpy as np
import onnxruntime as ort
TRUE = "rn1qkb1r/pb2pppp/5n2/1ppp4/8/3P1NP1/PPP1PPBP/RNBQ1RK1"
IDX2SYM = {0:"K",1:"Q",2:"R",3:"B",4:"N",5:"P",6:"k",7:"q",8:"r",9:"b",10:"n",11:"p"}
# ---- EXACT functions from Preprocess.ipynb ----
def crop_image(imagem):
cx, cy = imagem.shape[1] // 2, imagem.shape[0] // 2
return imagem[cy - 520:cy + 450, cx - 550:cx + 550]
def detect_edges(im): return cv.Canny(im, 100, 150, apertureSize=3)
def detect_lines(b): return cv.HoughLines(b, 1, np.pi / 180, 160)
def group_lines(linhas):
g = []
for linha in linhas:
rho, theta = linha[0]
if not any(abs(rho - l[0][0]) < 30 and abs(theta - l[0][1]) < np.pi / 18 for l in g):
g.append(linha)
return g
def filter_board_lines(g):
min_rho = None
for linha in g:
rho, theta = linha[0]
if 1 < theta < 3 and (min_rho is None or rho < min_rho):
min_rho = rho
return [l for l in g if l[0][0] != min_rho]
def compute_intersections(linhas_casas):
lp = []
for linha in linhas_casas:
rho, theta = linha[0]
a, b = np.cos(theta), np.sin(theta)
x0, y0 = a * rho, b * rho
lp.append([(int(x0 + 10000 * -b), int(y0 + 10000 * a)),
(int(x0 - 10000 * -b), int(y0 - 10000 * a)), theta])
rows = []
for i in range(len(lp)):
pts = []
for j in range(len(lp)):
x1, y1 = lp[i][0]; x2, y2 = lp[i][1]
x3, y3 = lp[j][0]; x4, y4 = lp[j][1]
det = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
if det != 0:
ix = int(((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / det)
iy = int(((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / det)
if 0 <= ix <= 1920 and 0 <= iy <= 1080 and 1 < lp[i][2] < 3:
pts.append((ix, iy))
if pts:
pts.sort(key=lambda p: p[0]); rows.append(pts)
rows.sort(key=lambda p: p[0][0])
return rows
def build_squares(rows):
casas = []
for i in range(len(rows) - 1):
for j in range(len(rows[i]) - 1):
casas.append([rows[i][j], rows[i][j + 1], rows[i + 1][j], rows[i + 1][j + 1]])
return casas
def crop_square(imagem, p):
xs = [pt[0] for pt in p]; ys = [pt[1] for pt in p]
x_min, x_max = min(xs), max(xs); y_min, y_max = min(ys), max(ys)
if y_min - 80 < 0:
y_min = 80
return imagem[y_min - 90:y_max, x_min:x_max]
def main():
img = cv.imread("samples/board.png")
im = crop_image(img)
print("cropped:", im.shape)
lines = detect_lines(detect_edges(im))
print("hough lines:", 0 if lines is None else len(lines))
g = group_lines(lines)
casas = build_squares(compute_intersections(filter_board_lines(g)))
print("squares detected:", len(casas))
# montage of crops in detected order
tiles = []
for p in casas[:64]:
c = crop_square(im, p)
tiles.append(cv.resize(c, (90, 90)) if c.size else np.zeros((90, 90, 3), np.uint8))
while len(tiles) < 64:
tiles.append(np.zeros((90, 90, 3), np.uint8))
mont = np.vstack([np.hstack(tiles[r * 8:r * 8 + 8]) for r in range(8)])
cv.imwrite("/tmp/sq_montage.png", mont)
print("saved /tmp/sq_montage.png")
if len(casas) != 64:
print("!! not 64 squares — pipeline needs tuning for this image"); return
sess = ort.InferenceSession("models_onnx/digitizer.fp32.onnx", providers=["CPUExecutionProvider"])
for nm, fn in [("/255", lambda r: r / 255.0),
("imagenet", lambda r: (r / 255.0 - [0.485,0.456,0.406]) / [0.229,0.224,0.225])]:
import chess
board = chess.Board(None)
for i, p in enumerate(casas[:64]):
c = crop_square(im, p)
if c.size == 0:
continue
rgb = cv.cvtColor(cv.resize(c, (96, 96)), cv.COLOR_BGR2RGB).astype(np.float32)
x = np.asarray(fn(rgb), np.float32).transpose(2, 0, 1)[None]
out = sess.run(None, {"input": x})[0].flatten()
board.set_piece_at(i, chess.Piece.from_symbol(IDX2SYM[int(out.argmax())]))
print(f"\n[{nm}] FEN: {board.board_fen()}")
print(f" TRUE: {TRUE}")
if __name__ == "__main__":
main()
|