| """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]): |
| 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() |
|
|
