Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Decisive test: 4-corner grid (robust) + EXACT Preprocess crop (bbox + 90px top), | |
| classify the 64 squares, compare to FEN under all 8 orientations and 3 norms.""" | |
| 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"} | |
| # corners on the FULL photo (TL,TR,BR,BL of the 8x8 area) | |
| CORNERS_FULL = np.array([[520,140],[1360,150],[1430,915],[452,918]], np.float32) | |
| def crop_image(im): | |
| cx, cy = im.shape[1]//2, im.shape[0]//2 | |
| return im[cy-520:cy+450, cx-550:cx+550], (cx-550, cy-520) | |
| def true_grid(): | |
| g=[] | |
| for row in TRUE.split("/"): | |
| r=[] | |
| for ch in row: r += ["."]*int(ch) if ch.isdigit() else [ch] | |
| g.append(r) | |
| return np.array(g) | |
| def main(): | |
| img = cv.imread("samples/board.png") | |
| im, (ox, oy) = crop_image(img) | |
| corners = CORNERS_FULL - [ox, oy] # into crop space | |
| unit = np.array([[0,0],[8,0],[8,8],[0,8]], np.float32) | |
| Hmat = cv.getPerspectiveTransform(unit, corners) # unit grid -> image | |
| def node(c, r): | |
| p = cv.perspectiveTransform(np.array([[[c, r]]], np.float32), Hmat)[0][0] | |
| return p | |
| crops = [] | |
| for r in range(8): | |
| for c in range(8): | |
| pts = [node(c, r), node(c+1, r), node(c+1, r+1), node(c, r+1)] | |
| xs = [p[0] for p in pts]; ys = [p[1] for p in pts] | |
| x1, x2 = int(min(xs)), int(max(xs)) | |
| y1 = max(0, int(min(ys)) - 90); y2 = int(max(ys)) | |
| crops.append(im[y1:y2, x1:x2]) | |
| sess = ort.InferenceSession("models_onnx/digitizer.fp32.onnx", providers=["CPUExecutionProvider"]) | |
| T = true_grid() | |
| occ = (T != ".") | |
| n_occ = int(occ.sum()) | |
| 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]), | |
| ("[-1,1]", lambda r: r/127.5-1)]: | |
| P = np.empty((8,8), object) | |
| for i, crop in enumerate(crops): | |
| if crop.size == 0: | |
| P[i//8, i%8] = "?"; continue | |
| rgb = cv.cvtColor(cv.resize(crop,(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() | |
| P[i//8, i%8] = IDX2SYM[int(out.argmax())] | |
| # best over 8 dihedral orientations, counting only TRUE-occupied squares | |
| best = 0; bestname = "" | |
| for k, Q in [("id",P),("rot90",np.rot90(P)),("rot180",np.rot90(P,2)),("rot270",np.rot90(P,3)), | |
| ("flipUD",np.flipud(P)),("flipLR",np.fliplr(P)), | |
| ("transp",P.T),("anti",np.fliplr(np.rot90(P)))]: | |
| if Q.shape!=T.shape: continue | |
| m = int(((Q==T)&occ).sum()) | |
| if m>best: best, bestname = m, k | |
| print(f"[{nm:9s}] best occupied-match {best:2d}/{n_occ} ({bestname})") | |
| if __name__ == "__main__": | |
| main() | |