Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Full image -> FEN: exact grid + crop + transform + empty detection (binarize | |
| center) + flipUD orientation. Compares to the known FEN for i10.""" | |
| import sys | |
| import cv2 as cv | |
| import numpy as np | |
| import onnxruntime as ort | |
| from digitize_full import get_casas, crop_square | |
| TRUE = "rn1qkb1r/pb2pppp/5n2/1ppp4/8/3P1NP1/PPP1PPBP/RNBQ1RK1" | |
| IDX2SYM = {0:"P",1:"N",2:"B",3:"R",4:"Q",5:"K",6:"p",7:"n",8:"b",9:"r",10:"q",11:"k"} | |
| MEAN, STD = np.array([0.485,0.456,0.406],np.float32), np.array([0.229,0.224,0.225],np.float32) | |
| def preprocess(bgr): | |
| rgb = cv.cvtColor(cv.resize(bgr,(224,224)),cv.COLOR_BGR2RGB).astype(np.float32)/255.0 | |
| return ((rgb-MEAN)/STD).transpose(2,0,1)[None] | |
| def is_empty(im, quad): | |
| cx = int(np.mean([p[0] for p in quad])); cy = int(np.mean([p[1] for p in quad])) | |
| w = max(6, int(0.22*(max(p[0] for p in quad)-min(p[0] for p in quad)))) | |
| patch = im[max(0,cy-w):cy+w, max(0,cx-w):cx+w] | |
| if patch.size == 0: | |
| return True | |
| g = cv.cvtColor(patch, cv.COLOR_BGR2GRAY) | |
| return g.std() < 18 # flat colour -> empty | |
| def grid_to_fen(grid): | |
| out=[] | |
| for row in grid: | |
| s=""; e=0 | |
| for cell in row: | |
| if cell==".": e+=1 | |
| else: | |
| if e: s+=str(e); e=0 | |
| s+=cell | |
| if e: s+=str(e) | |
| out.append(s) | |
| return "/".join(out) | |
| def main(): | |
| path = sys.argv[1] if len(sys.argv) > 1 else "samples/board.png" | |
| im, casas = get_casas(path) | |
| sess = ort.InferenceSession("models_onnx/digitizer3d.fp32.onnx", providers=["CPUExecutionProvider"]) | |
| syms=[] | |
| for p in casas[:64]: | |
| if is_empty(im, p): | |
| syms.append("."); continue | |
| c = crop_square(im, p) | |
| out = sess.run(None, {"input": preprocess(c)})[0].flatten() | |
| syms.append(IDX2SYM[int(out.argmax())]) | |
| grid = np.flipud(np.array(syms, object).reshape(8,8)) # flipUD orientation | |
| fen = grid_to_fen(grid) | |
| print("PRED FEN:", fen) | |
| print("TRUE FEN:", TRUE) | |
| T = np.array([list(r) for r in | |
| ["".join("." if ch.isdigit() and False else ch for ch in row) for row in TRUE.split("/")]], object) | |
| # build true grid properly | |
| tg=[] | |
| for row in TRUE.split("/"): | |
| rr=[] | |
| for ch in row: rr += ["."]*int(ch) if ch.isdigit() else [ch] | |
| tg.append(rr) | |
| T=np.array(tg,object) | |
| print(f"exact-cell match: {int((grid==T).sum())}/64") | |
| if __name__ == "__main__": | |
| main() | |