EasyMocap/easymocap/annotator/chessboard.py
2021-04-14 15:22:51 +08:00

38 lines
1.4 KiB
Python

import numpy as np
import cv2
def _findChessboardCorners(img, pattern):
"basic function"
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
retval, corners = cv2.findChessboardCorners(img, pattern,
flags=cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_FILTER_QUADS)
if not retval:
return False, None
corners = cv2.cornerSubPix(img, corners, (11, 11), (-1, -1), criteria)
corners = corners.squeeze()
return True, corners
def _findChessboardCornersAdapt(img, pattern):
"Adapt mode"
img = cv2.adaptiveThreshold(img, 255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\
cv2.THRESH_BINARY,21, 2)
return _findChessboardCorners(img, pattern)
def findChessboardCorners(img, annots, pattern):
if annots['visited']:
return None
annots['visited'] = True
gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
# Find the chess board corners
for func in [_findChessboardCornersAdapt, _findChessboardCorners]:
ret, corners = func(gray, pattern)
if ret:break
else:
return None
# found the corners
show = img.copy()
show = cv2.drawChessboardCorners(show, pattern, corners, ret)
assert corners.shape[0] == len(annots['keypoints2d'])
corners = np.hstack((corners, np.ones((corners.shape[0], 1))))
annots['keypoints2d'] = corners.tolist()
return show