update detect chessboard
This commit is contained in:
sq-titanv
parent
f8d551bb5e
commit
076d2fa834
@ -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)
|
||||
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)
|
||||
Loading…
Reference in New Issue
Block a user