digitize_full.py

#!/usr/bin/env python3
"""Full digitization on samples/board.png (i10): exact Preprocess grid (with the
3rd-from-bottom line removed) + exact bbox+90 crop + classify, vs known FEN."""
import cv2 as cv
import numpy as np
import onnxruntime as ort
import hough_pipeline as hp

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"}


def get_casas(path):
    im = hp.crop_image(cv.imread(path))
    board = hp.filter_board_lines(hp.group_lines(hp.detect_lines(hp.detect_edges(im))))
    rows = hp.compute_intersections(board)
    by_y = sorted(rows, key=lambda r: np.mean([p[1] for p in r]), reverse=True)
    rows = [r for r in rows if r is not by_y[2]]
    return im, hp.build_squares(rows)


def crop_square(im, 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)
    return im[max(0, y_min - 90):y_max, x_min:x_max]


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():
    im, casas = get_casas("samples/board.png")
    print("squares:", len(casas))
    if len(casas) != 64:
        return
    crops = [crop_square(im, p) for p in casas[:64]]
    sess = ort.InferenceSession("models_onnx/digitizer.fp32.onnx", providers=["CPUExecutionProvider"])
    T = true_grid(); occ = (T != "."); n = 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)]:
        syms = []
        for c in crops:
            if c.size == 0:
                syms.append("?"); 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()
            syms.append(IDX2SYM[int(out.argmax())])
        P = np.array(syms, object).reshape(8, 8)
        best = max(
            (int(((Q == T) & occ).sum()), k)
            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)), ("T", P.T), ("anti", np.fliplr(np.rot90(P)))]
        )
        print(f"[{nm:9s}] best occupied-match {best[0]:2d}/{n}  ({best[1]})")


if __name__ == "__main__":
    main()