File size: 2,405 Bytes
e8c58bd | 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 | """Find a scoreboard's fields and corners in a picture.
usage: python example.py PICTURE (needs onnxruntime, opencv-python, numpy)"""
import sys
from pathlib import Path
import cv2
import numpy as np
import onnxruntime as ort
HERE = Path(__file__).parent
CLASSES = ["home", "away", "clock", "period", "other", "board"]
CROP, MARGIN = 256, 0.15
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD = np.array([0.229, 0.224, 0.225], np.float32)
def detect(sess, bgr):
h, w = bgr.shape[:2]
rgb = cv2.cvtColor(cv2.resize(bgr, (640, 640), interpolation=cv2.INTER_LINEAR), cv2.COLOR_BGR2RGB)
x = (rgb.astype(np.float32) / 255).transpose(2, 0, 1)[None]
labels, boxes, scores = sess.run(None, {"images": x, "orig_target_sizes": np.array([[w, h]], np.int64)})
return labels[0], boxes[0], scores[0]
def corners(sess, bgr, box):
"""The board's four corners (top left, top right, bottom right, bottom left) in picture pixels."""
x1, y1, x2, y2 = box
mw, mh = MARGIN * (x2 - x1), MARGIN * (y2 - y1)
src = np.array([[x1 - mw, y1 - mh], [x2 + mw, y1 - mh], [x2 + mw, y2 + mh], [x1 - mw, y2 + mh]], np.float32)
H = cv2.getPerspectiveTransform(src, np.array([[0, 0], [CROP, 0], [CROP, CROP], [0, CROP]], np.float32))
crop = cv2.cvtColor(cv2.warpPerspective(bgr, H, (CROP, CROP)), cv2.COLOR_BGR2RGB).astype(np.float32) / 255
c = sess.run(None, {"crop": ((crop - MEAN) / STD).transpose(2, 0, 1)[None]})[0][0] * CROP
return cv2.perspectiveTransform(c.reshape(-1, 1, 2), np.linalg.inv(H)).reshape(4, 2)
def main():
bgr = cv2.imread(sys.argv[1])
fields = ort.InferenceSession(str(HERE / "fields_board.onnx"), providers=["CPUExecutionProvider"])
corner_net = ort.InferenceSession(str(HERE / "corners.onnx"), providers=["CPUExecutionProvider"])
labels, boxes, scores = detect(fields, bgr)
for i, name in enumerate(CLASSES[:4]): # the best box of each role
sel = (labels == i) & (scores >= 0.3)
if sel.any():
k = np.argmax(np.where(sel, scores, -1))
print(f"{name}: box {np.round(boxes[k]).astype(int).tolist()} score {scores[k]:.2f}")
sel = (labels == CLASSES.index("board")) & (scores >= 0.3)
if sel.any():
k = np.argmax(np.where(sel, scores, -1))
print("board corners:", np.round(corners(corner_net, bgr, boxes[k])).astype(int).tolist())
if __name__ == "__main__":
main()
|