186 lines
7.1 KiB
Python
186 lines
7.1 KiB
Python
import os
|
||
import numpy as np
|
||
import cv2 as cv
|
||
import glob
|
||
import os.path as osp
|
||
import json
|
||
from tqdm import tqdm
|
||
|
||
|
||
# 先想清楚文件夹结构
|
||
# extri文件夹：存放棋盘格照片，命名规则是cam1.jpg, cam2.jpg, cam3.jpg, ...
|
||
# 另一种模式是只检测2d点，不生成3d点，需要指定文件夹和输出路径
|
||
|
||
|
||
def write_json(data, output_path):
|
||
with open(output_path, "w") as f:
|
||
json.dump(data, f)
|
||
|
||
|
||
def read_json(input):
|
||
with open(input, "r") as f:
|
||
data = json.load(f)
|
||
return data
|
||
|
||
|
||
def read_img_paths(imgFolder):
|
||
imgPaths = []
|
||
for extension in ["jpg", "png", "jpeg", "bmp"]:
|
||
imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
|
||
return imgPaths
|
||
|
||
|
||
def create_output_folder(baseFolder, outputFolder):
|
||
folder = osp.join(baseFolder, outputFolder)
|
||
if not osp.exists(folder):
|
||
os.makedirs(folder)
|
||
return folder
|
||
|
||
base_path = "data"
|
||
extri_img_path = osp.join(base_path, "chessboard", "extri")
|
||
extri_vis_path = osp.join(base_path, "vis", "extri")
|
||
json_output_path = osp.join(base_path, 'output_json')
|
||
|
||
def _findChessboardCorners(img, pattern):
|
||
"basic function"
|
||
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
|
||
retval, corners = cv.findChessboardCorners(img, pattern,
|
||
flags=cv.CALIB_CB_ADAPTIVE_THRESH + cv.CALIB_CB_FAST_CHECK + cv.CALIB_CB_FILTER_QUADS)
|
||
if not retval:
|
||
return False, None
|
||
corners = cv.cornerSubPix(img, corners, (11, 11), (-1, -1), criteria)
|
||
corners = corners.squeeze()
|
||
return True, corners
|
||
|
||
|
||
def _findChessboardCornersAdapt(img, pattern):
|
||
"Adapt mode"
|
||
img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, \
|
||
cv.THRESH_BINARY, 21, 2)
|
||
return _findChessboardCorners(img, pattern)
|
||
|
||
# 检测棋盘格角点并且可视化
|
||
def findChessboardCorners(img_path, pattern, show=False):
|
||
img = cv.imread(img_path)
|
||
if img is None:
|
||
raise FileNotFoundError(f"Image not found at {img_path}")
|
||
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
||
|
||
# Find the chess board corners
|
||
for func in [_findChessboardCorners, _findChessboardCornersAdapt]:
|
||
ret, corners = func(gray, pattern)
|
||
if ret: break
|
||
else:
|
||
return None
|
||
# 附加置信度 1.0 并返回
|
||
kpts2d = np.hstack([corners, np.ones((corners.shape[0], 1))])
|
||
|
||
if show:
|
||
# Draw and display the corners
|
||
img_with_corners = cv.drawChessboardCorners(img, pattern, corners, ret)
|
||
# 标出棋盘格的原点
|
||
origin = tuple(corners[0].astype(int)) # 原点的像素坐标
|
||
cv.circle(img_with_corners, origin, 10, (0, 0, 255), -1) # 绘制原点
|
||
cv.putText(img_with_corners, "Origin", (origin[0] + 10, origin[1] - 10),
|
||
cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
|
||
|
||
# 标出最后一个点
|
||
last_point = tuple(corners[-1].astype(int)) # 角点数组的最后一个点
|
||
cv.circle(img_with_corners, last_point, 10, (0, 255, 0), -1) # 绿色圆点
|
||
cv.putText(img_with_corners, "Last", (last_point[0] + 15, last_point[1] - 15),
|
||
cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) # 添加文字 "Last"
|
||
|
||
# 显示图像
|
||
cv.imwrite(osp.join(extri_vis_path, osp.basename(img_path)), img_with_corners)
|
||
return kpts2d
|
||
|
||
|
||
# 根据棋盘格生成三维坐标，棋盘格坐标系原点在左上角（同时也是全局坐标原点）
|
||
# 设定标定板z轴朝上，yx表示棋盘在yx平面上
|
||
# easymocap
|
||
# 注意，采用11x8的棋盘格，长边是y轴11，短边是x轴8，可以用opencv试一下
|
||
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) # 3d坐标，默认向上的坐标轴为0
|
||
# 长边是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
|
||
|
||
|
||
# 检测文件夹下的所有棋盘格图片，生成3d点和2d点，存入json文件
|
||
# 图片应该按照cam0.jpg, cam1.jpg, cam2.jpg, ...的命名方式，要和内参文件夹对应
|
||
def detect_chessboard(pattern, gridSize):
|
||
imgPaths = read_img_paths(extri_img_path)
|
||
if len(imgPaths) == 0:
|
||
print("No images found!")
|
||
return
|
||
|
||
data = {}
|
||
for imgPath in tqdm(imgPaths):
|
||
camname = osp.basename(imgPath).split(".")[0]
|
||
keypoints2d = findChessboardCorners(imgPath, pattern, show=True)
|
||
if keypoints2d is not None:
|
||
keypoints3d = getChessboard3d(pattern, gridSize)
|
||
data[camname] = {
|
||
"keypoints2d": keypoints2d.tolist(),
|
||
"keypoints3d": keypoints3d.tolist(),
|
||
"pattern": pattern,
|
||
"gridSize": gridSize
|
||
}
|
||
json_path = osp.join(json_output_path, "chessboard_keypoints.json")
|
||
write_json(data, json_path)
|
||
print(f"Saved keypoints to {json_path}")
|
||
|
||
|
||
# 只检测2d点存入json文件
|
||
def detect_chessboard_2d(imgFolder, pattern, outJsonPath):
|
||
pass
|
||
|
||
|
||
def test_findChessboardCorners(img_path, pattern, saveDir):
|
||
imgpaths = read_img_paths(img_path)
|
||
for imgpath in imgpaths:
|
||
img = cv.imread(imgpath)
|
||
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
||
ret, corners = cv.findChessboardCorners(gray, pattern)
|
||
if ret:
|
||
# 在棋盘格上绘制角点
|
||
img_with_corners = cv.drawChessboardCorners(img, pattern, corners, ret)
|
||
|
||
# 标出原点
|
||
origin = tuple(corners[0][0]) # 角点数组的第一个点作为原点
|
||
cv.circle(img_with_corners, (int(origin[0]), int(origin[1])), 10, (0, 0, 255), -1) # 红色圆点
|
||
cv.putText(img_with_corners, "Origin", (int(origin[0]) + 15, int(origin[1]) - 15),
|
||
cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) # 添加文字 "Origin"
|
||
|
||
# 标出最后一个点
|
||
last_point = tuple(corners[-1][0]) # 角点数组的最后一个点
|
||
cv.circle(img_with_corners, (int(last_point[0]), int(last_point[1])), 10, (0, 255, 0), -1) # 绿色圆点
|
||
cv.putText(img_with_corners, "Last", (int(last_point[0]) + 15, int(last_point[1]) - 15),
|
||
cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) # 添加文字 "Last"
|
||
|
||
# 保存带角点的图像
|
||
cv.imwrite(osp.join(saveDir, osp.basename(imgpath)), img_with_corners)
|
||
else:
|
||
print(f"Failed to detect chessboard corners in {imgpath}")
|
||
print(f"Saved images to {saveDir}")
|
||
|
||
if __name__ == '__main__':
|
||
# test1
|
||
img_path = "data/chessboard/extri"
|
||
pattern = (11, 8)
|
||
# saveDir = "data/test1"
|
||
# os.makedirs(saveDir, exist_ok=True)
|
||
# test_findChessboardCorners(img_path, pattern, saveDir)
|
||
detect_chessboard(pattern, 60)
|
||
|