test.py - randomcrap - random crap programs of varying quality
(HTM) git clone git://git.codemadness.org/randomcrap
(DIR) Log
(DIR) Files
(DIR) Refs
(DIR) README
(DIR) LICENSE
---
test.py (1840B)
---
1 # experiment for chess: simple detecting of pieces from board image.
2 # testing different match templates.
3 # it doesn't do any scaling/transforms so the board size and image to detect
4 # are unscaled etc.
5
6 import cv2
7 import numpy as np
8
9 # Load the chess board and chess piece images
10 img_board = cv2.imread('board.png')
11 img_piece = cv2.imread('piece.png')
12
13 #img_piece_mask_bg_img = cv2.imread('piece_mask.png')
14 #img_piece_mask_bg = cv2.cvtColor(img_piece_mask_bg_img, cv2.COLOR_BGR2GRAY)
15
16 img_piece_mask = cv2.cvtColor(img_piece, cv2.COLOR_BGR2GRAY)
17
18 # Convert both images to grayscale
19 img_board_gray = cv2.cvtColor(img_board, cv2.COLOR_BGR2GRAY)
20
21 # Apply morphological operations to extract the chess piece from the board
22 kernel = np.ones((3, 3), np.uint8)
23 img_piece_mask = cv2.erode(img_piece_mask, kernel, iterations=1)
24 img_piece_mask = cv2.dilate(img_piece_mask, kernel, iterations=1)
25
26 # https://docs.opencv.org/3.4/d4/dc6/tutorial_py_template_matching.html
27
28 w, h = img_piece_mask.shape[::-1]
29
30 #ret, mask = cv2.threshold(img_piece_mask, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
31 #cv2.imwrite('/tmp/mask.png', mask) # DEBUG
32 #cv2.imwrite('/tmp/piecemask.png', img_piece_mask) # DEBUG
33
34 # Find the matching location on the board
35 #res = cv2.matchTemplate(img_board_gray, img_piece_mask, method=cv2.TM_SQDIFF, mask=mask)
36 #threshold = 8
37
38 res = cv2.matchTemplate(img_board_gray, img_piece_mask, method=cv2.TM_CCOEFF_NORMED, mask=img_piece_mask)
39 threshold = 0.8
40
41 #res = cv2.matchTemplate(img_board_gray, img_piece_mask, method=cv2.TM_CCORR_NORMED, mask=mask)
42 #threshold = 0.8
43
44 loc = np.where(res >= threshold)
45 #loc = np.where(res < threshold) # cv2.TM_SQDIFF
46 for pt in zip(*loc[::-1]):
47 cv2.rectangle(img_board, pt, (pt[0] + w, pt[1] + h), (0,0,255), 2)
48
49 # Show the result
50 cv2.imshow('Result', img_board)
51 cv2.waitKey(0)
52 cv2.destroyAllWindows()