diff --git a/apps/calibration/detect_chessboard.py b/apps/calibration/detect_chessboard.py index ad45cd4..e9bb017 100644 --- a/apps/calibration/detect_chessboard.py +++ b/apps/calibration/detect_chessboard.py @@ -2,113 +2,193 @@ @ Date: 2021-07-16 20:13:57 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-07-21 19:56:38 - @ FilePath: /EasyMocap/apps/calibration/detect_chessboard.py + @ LastEditTime: 2022-05-11 20:41:10 + @ FilePath: /EasyMocapPublic/apps/calibration/detect_chessboard.py ''' # detect the corner of chessboard from easymocap.annotator.file_utils import getFileList, read_json, save_json +from easymocap.mytools.debug_utils import mywarn from tqdm import tqdm from easymocap.annotator import ImageFolder -from easymocap.annotator.chessboard import getChessboard3d, findChessboardCorners +from easymocap.annotator.chessboard import findChessboardCorners import numpy as np from os.path import join import cv2 import os +import func_timeout +import threading +from easymocap.mytools.debug_utils import log -def create_chessboard(path, pattern, gridSize, ext): +def getChessboard3d(pattern, gridSize, axis='yx'): + # 注意：这里为了让标定板z轴朝上，设定了短边是x，长边是y + template = np.mgrid[0:pattern[0], 0:pattern[1]].T.reshape(-1,2) + object_points = np.zeros((pattern[1]*pattern[0], 3), np.float32) + # 长边是x,短边是z + if axis == 'xz': + object_points[:, 0] = template[:, 0] + object_points[:, 2] = template[:, 1] + elif axis == 'yx': + object_points[:, 0] = template[:, 1] + object_points[:, 1] = template[:, 0] + else: + raise NotImplementedError + object_points = object_points * gridSize + return object_points + +def create_chessboard(path, image, pattern, gridSize, ext, overwrite=True): print('Create chessboard {}'.format(pattern)) - keypoints3d = getChessboard3d(pattern, gridSize=gridSize) + keypoints3d = getChessboard3d(pattern, gridSize=gridSize, axis=args.axis) keypoints2d = np.zeros((keypoints3d.shape[0], 3)) - imgnames = getFileList(path, ext=ext) + imgnames = getFileList(join(path, image), ext=ext) template = { 'keypoints3d': keypoints3d.tolist(), 'keypoints2d': keypoints2d.tolist(), + 'pattern': pattern, + 'grid_size': gridSize, 'visited': False } for imgname in tqdm(imgnames, desc='create template chessboard'): - annname = imgname.replace('images', 'chessboard').replace(ext, '.json') - annname = join(path, annname) - if os.path.exists(annname): + annname = imgname.replace(ext, '.json') + annname = join(path, 'chessboard', annname) + if os.path.exists(annname) and overwrite: # 覆盖keypoints3d data = read_json(annname) data['keypoints3d'] = template['keypoints3d'] save_json(annname, data) + elif os.path.exists(annname) and not overwrite: + continue else: save_json(annname, template) -def detect_chessboard(path, out, pattern, gridSize, args): - create_chessboard(path, pattern, gridSize, ext=args.ext) - dataset = ImageFolder(path, annot='chessboard', ext=args.ext) - dataset.isTmp = False - if args.silent: - trange = range(len(dataset)) - else: - trange = tqdm(range(len(dataset))) - for i in trange: - imgname, annotname = dataset[i] +def _detect_chessboard(datas, path, image, out, pattern): + for imgname, annotname in datas: + # imgname, annotname = dataset[i] # detect the 2d chessboard img = cv2.imread(imgname) annots = read_json(annotname) - show = findChessboardCorners(img, annots, pattern) + try: + show = findChessboardCorners(img, annots, pattern) + except func_timeout.exceptions.FunctionTimedOut: + show = None save_json(annotname, annots) if show is None: - if args.debug: - print('Cannot find {}'.format(imgname)) + mywarn('[Info] Cannot find chessboard in {}'.format(imgname)) continue - outname = join(out, imgname.replace(path + '/images/', '')) + outname = join(out, imgname.replace(path + '/{}/'.format(image), '')) os.makedirs(os.path.dirname(outname), exist_ok=True) - cv2.imwrite(outname, show) + if isinstance(show, np.ndarray): + cv2.imwrite(outname, show) -def detect_chessboard_sequence(path, out, pattern, gridSize, args): - create_chessboard(path, pattern, gridSize, ext=args.ext) - subs = sorted(os.listdir(join(path, 'images'))) +def detect_chessboard(path, image, out, pattern, gridSize, args): + create_chessboard(path, image, pattern, gridSize, ext=args.ext, overwrite=args.overwrite3d) + dataset = ImageFolder(path, image=image, annot='chessboard', ext=args.ext) + dataset.isTmp = False + trange = list(range(len(dataset))) + threads = [] + for i in range(args.mp): + ranges = trange[i::args.mp] + datas = [dataset[t] for t in ranges] + thread = threading.Thread(target=_detect_chessboard, args=(datas, path, image, out, pattern)) # 应该不存在任何数据竞争 + thread.start() + threads.append(thread) + for thread in threads: + thread.join() + + +def _detect_by_search(path, image, out, pattern, sub): + dataset = ImageFolder(path, sub=sub, annot='chessboard', ext=args.ext) + dataset.isTmp = False + nFrames = len(dataset) + found = np.zeros(nFrames, dtype=bool) + visited = np.zeros(nFrames, dtype=bool) + proposals = [] + init_step = args.max_step + min_step = args.min_step + for nf in range(0, nFrames, init_step): + if nf + init_step < len(dataset): + proposals.append([nf, nf+init_step]) + while len(proposals) > 0: + left, right = proposals.pop(0) + print('[detect] {} {:4.1f}% Check [{:5d}, {:5d}]'.format( + sub, visited.sum()/visited.shape[0]*100, left, right), end=' ') + for nf in [left, right]: + if not visited[nf]: + visited[nf] = True + imgname, annotname = dataset[nf] + # detect the 2d chessboard + img = cv2.imread(imgname) + annots = read_json(annotname) + try: + show = findChessboardCorners(img, annots, pattern) + except func_timeout.exceptions.FunctionTimedOut: + show = None + save_json(annotname, annots) + if show is None: + if args.debug: + print('[Info] Cannot find chessboard in {}'.format(imgname)) + found[nf] = False + continue + found[nf] = True + outname = join(out, imgname.replace(path + '{}{}{}'.format(os.sep, image, os.sep), '')) + os.makedirs(os.path.dirname(outname), exist_ok=True) + if isinstance(show, np.ndarray): + cv2.imwrite(outname, show) + print('{}-{}'.format('o' if found[left] else 'x', 'o' if found[right] else 'x')) + if not found[left] and not found[right]: + visited[left:right] = True + continue + mid = (left+right)//2 + if mid == left or mid == right: + continue + if mid - left > min_step: + proposals.append((left, mid)) + if right - mid > min_step: + proposals.append((mid, right)) + +def detect_chessboard_sequence(path, image, out, pattern, gridSize, args): + create_chessboard(path, image, pattern, gridSize, ext=args.ext, overwrite=args.overwrite3d) + subs = sorted(os.listdir(join(path, image))) + subs = [s for s in subs if os.path.isdir(join(path, image, s))] + if len(subs) == 0: + subs = [None] + from multiprocessing import Process + tasks = [] + for sub in subs: + task = Process(target=_detect_by_search, args=(path, image, out, pattern, sub)) + task.start() + tasks.append(task) + for task in tasks: + task.join() for sub in subs: dataset = ImageFolder(path, sub=sub, annot='chessboard', ext=args.ext) dataset.isTmp = False - nFrames = len(dataset) - found = np.zeros(nFrames, dtype=np.bool) - visited = np.zeros(nFrames, dtype=np.bool) - proposals = [] - init_step = args.max_step - min_step = args.min_step - for nf in range(0, nFrames, init_step): - if nf + init_step < len(dataset): - proposals.append([nf, nf+init_step]) - while len(proposals) > 0: - left, right = proposals.pop(0) - print('Check [{}, {}]'.format(left, right)) - for nf in [left, right]: - if not visited[nf]: - visited[nf] = True - imgname, annotname = dataset[nf] - # detect the 2d chessboard - img = cv2.imread(imgname) - annots = read_json(annotname) - show = findChessboardCorners(img, annots, pattern) - save_json(annotname, annots) - if show is None: - if args.debug: - print('Cannot find {}'.format(imgname)) - found[nf] = False - continue - found[nf] = True - outname = join(out, imgname.replace(path + '/images/', '')) - os.makedirs(os.path.dirname(outname), exist_ok=True) - cv2.imwrite(outname, show) - if not found[left] and not found[right]: - continue - mid = (left+right)//2 - if mid == left or mid == right: - continue - if mid - left > min_step: - proposals.append((left, mid)) - if right - mid > min_step: - proposals.append((mid, right)) + count, visited = 0, 0 + for nf in range(len(dataset)): + imgname, annotname = dataset[nf] + # detect the 2d chessboard + annots = read_json(annotname) + if annots['visited']: + visited += 1 + if annots['keypoints2d'][0][-1] > 0.01: + count += 1 + log('{}: found {:4d}/{:4d} frames'.format(sub, count, visited)) + +def check_chessboard(path, out): + subs_notvalid = [] + for sub in sorted(os.listdir(join(path, 'images'))): + if os.path.exists(join(out, sub)): + continue + subs_notvalid.append(sub) + print(subs_notvalid) + mywarn('Cannot find chessboard in view {}'.format(subs_notvalid)) + mywarn('Please annot them manually:') + mywarn(f'python3 apps/annotation/annot_calib.py {path} --mode chessboard --annot chessboard --sub {" ".join(subs_notvalid)}') if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument('path', type=str) + parser.add_argument('--image', type=str, default='images') parser.add_argument('--out', type=str, required=True) parser.add_argument('--ext', type=str, default='.jpg', choices=['.jpg', '.png']) parser.add_argument('--pattern', type=lambda x: (int(x.split(',')[0]), int(x.split(',')[1])), @@ -117,12 +197,20 @@ if __name__ == "__main__": help='The length of the grid size (unit: meter)') parser.add_argument('--max_step', type=int, default=50) parser.add_argument('--min_step', type=int, default=0) + parser.add_argument('--mp', type=int, default=4) + parser.add_argument('--axis', type=str, default='yx') parser.add_argument('--silent', action='store_true') parser.add_argument('--debug', action='store_true') + parser.add_argument('--overwrite3d', action='store_true') parser.add_argument('--seq', action='store_true') + parser.add_argument('--check', action='store_true') + args = parser.parse_args() if args.seq: - detect_chessboard_sequence(args.path, args.out, pattern=args.pattern, gridSize=args.grid, args=args) + detect_chessboard_sequence(args.path, args.image, args.out, pattern=args.pattern, gridSize=args.grid, args=args) else: - detect_chessboard(args.path, args.out, pattern=args.pattern, gridSize=args.grid, args=args) \ No newline at end of file + detect_chessboard(args.path, args.image, args.out, pattern=args.pattern, gridSize=args.grid, args=args) + + if args.check: + check_chessboard(args.path, args.out) \ No newline at end of file