EasyMocap/easymocap/annotator/chessboard.py

'''
  @ Date: 2021-04-13 16:14:36
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
  @ LastEditTime: 2022-10-25 20:56:26
  @ FilePath: /EasyMocapRelease/easymocap/annotator/chessboard.py
'''
import numpy as np
import cv2
from func_timeout import func_set_timeout

def getChessboard3d(pattern, gridSize, axis='xy'):
    object_points = np.zeros((pattern[1]*pattern[0], 3), np.float32)
    # 注意：这里为了让标定板z轴朝上，设定了短边是x，长边是y
    object_points[:,:2] = np.mgrid[0:pattern[0], 0:pattern[1]].T.reshape(-1,2)
    object_points[:, [0, 1]] = object_points[:, [1, 0]]
    object_points = object_points * gridSize
    if axis == 'zx':
        object_points = object_points[:, [1, 2, 0]]
    return object_points

colors_chessboard_bar = [
    [0, 0, 255],
    [0, 128, 255],
    [0, 200, 200],
    [0, 255, 0],
    [200, 200, 0],
    [255, 0, 0],
    [255, 0, 250]
]

def get_lines_chessboard(pattern=(9, 6)):
    w, h = pattern[0], pattern[1]
    lines = []
    lines_cols = []
    for i in range(w*h-1):
        lines.append([i, i+1])
        lines_cols.append(colors_chessboard_bar[(i//w)%len(colors_chessboard_bar)])
    return lines, lines_cols

def _findChessboardCorners(img, pattern, debug):
    "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, debug):
    "Adapt mode"
    img = cv2.adaptiveThreshold(img, 255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\
                cv2.THRESH_BINARY, 21, 2)
    return _findChessboardCorners(img, pattern, debug)

@func_set_timeout(5)
def findChessboardCorners(img, annots, pattern, debug=False):
    conf = sum([v[2] for v in annots['keypoints2d']])
    if annots['visited'] and conf > 0:
        return True
    elif annots['visited']:
        return None
    annots['visited'] = True
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    # Find the chess board corners
    for func in [_findChessboardCorners, _findChessboardCornersAdapt]:
        ret, corners = func(gray, pattern, debug)
        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

def create_chessboard(path, keypoints3d, out='annots'):
    from tqdm import tqdm
    from os.path import join
    from .file_utils import getFileList, save_json, read_json
    import os
    keypoints2d = np.zeros((keypoints3d.shape[0], 3))
    imgnames = getFileList(join(path, 'images'), ext='.jpg', max=1)
    imgnames = [join('images', i) for i in imgnames]
    template = {
        'keypoints3d': keypoints3d.tolist(),
        'keypoints2d': keypoints2d.tolist(),
        'visited': False
    }
    for imgname in tqdm(imgnames, desc='create template chessboard'):
        annname = imgname.replace('images', out).replace('.jpg', '.json')
        annname = join(path, annname)
        if not os.path.exists(annname):
            save_json(annname, template)
        elif True:
            annots = read_json(annname)
            annots['keypoints3d'] = template['keypoints3d']
            save_json(annname, annots)


def detect_charuco(image, aruco_type, long, short, squareLength, aruco_len):
    ARUCO_DICT = {
        "4X4_50": cv2.aruco.DICT_4X4_50,
        "4X4_100": cv2.aruco.DICT_4X4_100,
        "5X5_100": cv2.aruco.DICT_5X5_100,
        "5X5_250": cv2.aruco.DICT_5X5_250,
    }
    # 创建ChArUco标定板
    dictionary = cv2.aruco.getPredefinedDictionary(dict=ARUCO_DICT[aruco_type])
    board = cv2.aruco.CharucoBoard_create(
        squaresY=long,
        squaresX=short,
        squareLength=squareLength,
        markerLength=aruco_len,
        dictionary=dictionary,
    )
    corners = board.chessboardCorners
    # ATTN: exchange the XY
    corners3d = corners[:, [1, 0, 2]]
    keypoints2d = np.zeros_like(corners3d)
    # 查找标志块的左上角点
    corners, ids, _ = cv2.aruco.detectMarkers(
        image=image, dictionary=dictionary, parameters=None
    )
    # 棋盘格黑白块内角点
    if ids is not None:
        retval, charucoCorners, charucoIds = cv2.aruco.interpolateCornersCharuco(
            markerCorners=corners, markerIds=ids, image=image, board=board
        )
        if retval:
            ids = charucoIds[:, 0]
            pts = charucoCorners[:, 0]
            keypoints2d[ids, :2] = pts
            keypoints2d[ids, 2] = 1.
    else:
        retval = False
    return retval, keypoints2d, corners3d

class CharucoBoard:
    def __init__(self, long, short, squareLength, aruco_len, aruco_type) -> None:    
        '''
            short,long 分别表示短边、长边的格子数.
            squareLength,aruco_len 分别表示棋盘格的边长与aruco的边长.
            aruco_type 表示Aruco的类型 4X4表示aruco中的白色格子是4x4的 _50表示aruco字典中有多少种aruco.
        '''
        # 定义现有的Aruco类型
        self.ARUCO_DICT = {
            "4X4_50": cv2.aruco.DICT_4X4_50,
            "4X4_100": cv2.aruco.DICT_4X4_100,
            "5X5_100": cv2.aruco.DICT_5X5_100,
            "5X5_250": cv2.aruco.DICT_5X5_250,
        }
        # 创建ChArUco标定板
        dictionary = cv2.aruco.getPredefinedDictionary(dict=self.ARUCO_DICT[aruco_type])
        board = cv2.aruco.CharucoBoard_create(
            squaresY=long,
            squaresX=short,
            squareLength=squareLength,
            markerLength=aruco_len,
            dictionary=dictionary,
        )
        corners = board.chessboardCorners
        # ATTN: exchange the XY
        corners = corners[:, [1, 0, 2]]
        self.template = {
            'keypoints3d': corners,
            'keypoints2d': np.zeros_like(corners),
            'pattern': (long-1, short-1),
            'grid_size': squareLength,
            'visted': False
        }
        print(corners.shape)
        self.dictionary = dictionary
        self.board = board
    
    def detect(self, img_color, annots):
        # 查找标志块的左上角点
        corners, ids, _ = cv2.aruco.detectMarkers(
            image=img_color, dictionary=self.dictionary, parameters=None
        )
        # 棋盘格黑白块内角点
        if ids is not None:
            retval, charucoCorners, charucoIds = cv2.aruco.interpolateCornersCharuco(
                markerCorners=corners, markerIds=ids, image=img_color, board=self.board
            )
        else:
            retval = False
        if retval:
            # 绘制棋盘格黑白块内角点
            cv2.aruco.drawDetectedCornersCharuco(
                img_color, charucoCorners, charucoIds, [0, 0, 255]
            )
            if False:
                cv2.aruco.drawDetectedMarkers(
                    image=img_color, corners=corners, ids=ids, borderColor=None
                )

            ids = charucoIds[:, 0]
            pts = charucoCorners[:, 0]
            annots['keypoints2d'][ids, :2] = pts
            annots['keypoints2d'][ids, 2] = 1.
            # if args.show:
            #     img_color = cv2.resize(img_color, None, fx=0.5, fy=0.5)
            #     cv2.imshow('vis', img_color)
            #     cv2.waitKey(0)
            # visname = imgname.replace(images, output)
            # os.makedirs(os.path.dirname(visname), exist_ok=True)
            # cv2.imwrite(visname, img_color)
        else:
            # mywarn('Cannot find in {}'.format(imgname))
            pass
        
    def __call__(self, imgname, images='images', output='output'):
        import os
        from .file_utils import read_json, save_json
        import copy
        img_color = cv2.imread(imgname)
        annotname = imgname.replace('images', 'chessboard').replace('.jpg', '.json')
        if os.path.exists(annotname):
            annots = read_json(annotname)
            if annots['visited']:
                return
        else:
            annots = copy.deepcopy(self.template)
        annots['visited'] = True
        self.detect(img_color, annots)
        annots['keypoints2d'] = annots['keypoints2d'].tolist()
        annots['keypoints3d'] = annots['keypoints3d'].tolist()
        save_json(annotname, annots)