probe_norm.py

#!/usr/bin/env python3
"""Probe: is the digitizer responding at all? Try fp32 vs int8 and a few
normalizations on the 64 squares; report match-count and mean confidence."""
import cv2 as cv
import numpy as np
import onnxruntime as ort
import torch
from app.chess_models import PieceImageClassifier

SIZE, MARGIN = 800, 0.6
CORNERS = np.array([[520, 140], [1360, 150], [1430, 915], [452, 918]], np.float32)
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"}
IMEAN, ISTD = np.array([0.485,0.456,0.406],np.float32), np.array([0.229,0.224,0.225],np.float32)

# re-export an fp32 digitizer (no quantization) for comparison
m = PieceImageClassifier()
ck = torch.load("../Master-of-Science-Degree-Project/Models/imageClassifierReal.pth", map_location="cpu", weights_only=False)
sd = ck if all(isinstance(v, torch.Tensor) for v in ck.values()) else next(v for v in ck.values() if isinstance(v, dict))
m.load_state_dict(sd); m.eval()
torch.onnx.export(m, torch.randn(1,3,96,96), "models_onnx/digitizer.fp32.onnx",
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input":{0:"b",2:"h",3:"w"}}, opset_version=17, dynamo=False)

def grid_true():
    g=[]
    for row in TRUE.split("/"):
        r=[]
        for ch in row: r += ["."]*int(ch) if ch.isdigit() else [ch]
        g.append(r)
    return g

def crops(img):
    dst=np.array([[0,0],[SIZE,0],[SIZE,SIZE],[0,SIZE]],np.float32)
    Minv=cv.getPerspectiveTransform(dst,CORNERS); cell=SIZE/8; out=[]
    for r in range(8):
        for c in range(8):
            y0=(r-MARGIN)*cell
            q=np.array([[[c*cell,y0],[(c+1)*cell,y0],[(c+1)*cell,(r+1)*cell],[c*cell,(r+1)*cell]]],np.float32)
            o=cv.perspectiveTransform(q,Minv)[0]; xs,ys=o[:,0],o[:,1]
            x1,x2=max(0,int(xs.min())),min(img.shape[1],int(xs.max()))
            y1,y2=max(0,int(ys.min())),min(img.shape[0],int(ys.max()))
            out.append(img[y1:y2,x1:x2])
    return out

def norms(bgr):
    rgb=cv.cvtColor(cv.resize(bgr,(96,96)),cv.COLOR_BGR2RGB).astype(np.float32)
    return {
        "/255":            (rgb/255.0).transpose(2,0,1)[None],
        "imagenet":        ((rgb/255.0-IMEAN)/ISTD).transpose(2,0,1)[None],
        "raw255":          rgb.transpose(2,0,1)[None],
        "[-1,1]":          (rgb/127.5-1).transpose(2,0,1)[None],
    }

img=cv.imread("samples/board.png"); cs=crops(img); true=[s for row in grid_true() for s in row]
for model in ["digitizer.int8.onnx","digitizer.fp32.onnx"]:
    sess=ort.InferenceSession(f"models_onnx/{model}",providers=["CPUExecutionProvider"])
    for nm in ["/255","imagenet","raw255","[-1,1]"]:
        hits=0; probs=[]
        for crop,t in zip(cs,true):
            if crop.size==0: continue
            x=norms(crop)[nm]
            out=sess.run(None,{"input":x.astype(np.float32)})[0].flatten()
            p=np.exp(out)/np.exp(out).sum(); k=int(p.argmax()); probs.append(p[k])
            if IDX2SYM[k]==t: hits+=1
        print(f"{model:22s} {nm:9s} match {hits:2d}/64  meanconf {np.mean(probs):.2f}")