203 lines
7.5 KiB
Python
203 lines
7.5 KiB
Python
'''
|
|
@ Date: 2021-05-26 13:19:12
|
|
@ Author: Qing Shuai
|
|
@ LastEditors: Qing Shuai
|
|
@ LastEditTime: 2022-09-26 20:13:53
|
|
@ FilePath: /EasyMocapPublic/apps/annotation/annot_calib.py
|
|
'''
|
|
from easymocap.annotator.basic_visualize import plot_text, resize_to_screen, vis_bbox, vis_line
|
|
from easymocap.mytools.debug_utils import mywarn
|
|
from easymocap.mytools.vis_base import plot_point
|
|
from easymocap.annotator import ImageFolder
|
|
from easymocap.annotator import vis_point
|
|
from easymocap.annotator import AnnotBase
|
|
from easymocap.mytools import read_json, save_json
|
|
from easymocap.mytools import plot_cross, plot_line, get_rgb
|
|
import numpy as np
|
|
import cv2
|
|
from tqdm import tqdm
|
|
from os.path import join
|
|
import os
|
|
from easymocap.annotator.chessboard import get_lines_chessboard, colors_chessboard_bar, create_chessboard
|
|
from easymocap.annotator.vanish_callback import calc_vanishpoint
|
|
|
|
class Matcher:
|
|
def __init__(self, path, mode, args) -> None:
|
|
if mode == 'chessboard':
|
|
pattern = args.pattern
|
|
lines, lines_color = get_lines_chessboard(pattern)
|
|
self.nJoints = pattern[0]*pattern[1]
|
|
else:
|
|
annots = read_json(join(path, 'calib.json'))
|
|
lines = annots['lines']
|
|
if 'lines_color' in annots.keys():
|
|
lines_color = annots['lines_color']
|
|
else:
|
|
lines_color = [colors_chessboard_bar[i%len(colors_chessboard_bar)] for i in range(len(lines))]
|
|
keypoints3d = np.array(annots['keypoints3d'])
|
|
create_chessboard(path, keypoints3d, out=args.annot)
|
|
self.nJoints = len(keypoints3d)
|
|
self.lines = lines
|
|
self.lines_color = lines_color
|
|
self.cache_lines = []
|
|
self.cnt = -1
|
|
self.hint()
|
|
|
|
def hint(self):
|
|
self.cnt = (self.cnt + 1)%self.nJoints
|
|
print('>>> label point {}'.format(self.cnt))
|
|
|
|
def back(self, annotator, param, conf=1.):
|
|
"switch to previous points"
|
|
self.cnt -= 2
|
|
self.hint()
|
|
|
|
def add(self, annotator, param, conf=1.):
|
|
"switch to next points"
|
|
click = param['click']
|
|
if click is not None:
|
|
param['annots']['keypoints2d'][self.cnt] = [click[0], click[1], conf]
|
|
param['annots']['visited'] = True
|
|
param['click'] = None
|
|
self.hint()
|
|
|
|
def add_conf(self, annotator, param):
|
|
self.add(annotator, param, conf=0.5)
|
|
|
|
def add_point_by_2lines(self, annotator, param):
|
|
start = param['start']
|
|
end = param['end']
|
|
if start is None:
|
|
return 0
|
|
if len(self.cache_lines) < 2:
|
|
self.cache_lines.append((start, end))
|
|
param['start'] = None
|
|
param['end'] = None
|
|
if len(self.cache_lines) == 2:
|
|
# calculate intersect
|
|
inp = np.zeros((2, 2, 3)) # 2, points, (x, y, c)
|
|
for i in range(len(self.cache_lines)):
|
|
start, end = self.cache_lines[i]
|
|
inp[0, i, 0] = start[0]
|
|
inp[0, i, 1] = start[1]
|
|
inp[0, i, 2] = 1.
|
|
inp[1, i, 0] = end[0]
|
|
inp[1, i, 1] = end[1]
|
|
inp[1, i, 2] = 1.
|
|
intersect = calc_vanishpoint(inp)
|
|
click = (int(intersect[0]), int(intersect[1]))
|
|
param['click'] = click
|
|
self.cache_lines = []
|
|
|
|
def clear(self, annotator, param):
|
|
"clear all points"
|
|
for i in range(self.nJoints):
|
|
param['annots']['keypoints2d'][i][2] = 0.
|
|
self.cnt = self.nJoints - 1
|
|
self.hint()
|
|
|
|
def clear_point(self, annotator, param):
|
|
"clear current points"
|
|
param['annots']['keypoints2d'][self.cnt][2] = 0.
|
|
|
|
def vis(self, img, annots, **kwargs):
|
|
border = 100
|
|
text_size = 2
|
|
width = 5
|
|
cv2.putText(img, 'Current {}: {:.0f}'.format(self.cnt, annots['keypoints2d'][self.cnt][2]), (border, border), cv2.FONT_HERSHEY_SIMPLEX, text_size, (0, 0, 255), width)
|
|
if kwargs['click'] is not None:
|
|
return img
|
|
lw = max(int(round(img.shape[0]/500)), 1)
|
|
width = lw * 5
|
|
k2d = np.array(annots['keypoints2d'])
|
|
# cache lines
|
|
for nl, (start, end) in enumerate(self.cache_lines):
|
|
plot_line(img, start, end, lw, (255, 255, 200))
|
|
|
|
for nl, (i, j) in enumerate(self.lines):
|
|
if k2d[i][2] > 0 and k2d[j][2] > 0:
|
|
plot_line(img, k2d[i], k2d[j], lw, self.lines_color[nl])
|
|
for i, (x, y, c) in enumerate(k2d):
|
|
if c > 0:
|
|
plot_cross(img, x, y, self.lines_color[min(len(self.lines)-1, i)], width=width, lw=lw)
|
|
plot_point(img, x, y, r=lw*2, col=self.lines_color[min(len(self.lines)-1, i)], pid=i)
|
|
if i == self.cnt:
|
|
plot_point(img, x, y, r=lw*16, col=(127, 127, 255), pid=-1, circle_type=lw*2)
|
|
return img
|
|
|
|
def print(self, annotator, **kwargs):
|
|
print(self.annots)
|
|
|
|
def detect(self, annotator, param):
|
|
"detect chessboard"
|
|
start = param['start']
|
|
end = param['end']
|
|
if start is None:
|
|
return 0
|
|
import cv2
|
|
crop = param['img0'][start[1]:end[1], start[0]:end[0]]
|
|
from easymocap.annotator.chessboard import findChessboardCorners
|
|
pattern = args.pattern
|
|
annots = {'visited':False, 'keypoints2d':np.zeros((pattern[0]*pattern[1], 3))}
|
|
print('Redetect the chessboard...')
|
|
if False:
|
|
show = findChessboardCorners(crop, annots, args.pattern, debug=True)
|
|
else:
|
|
self.detect_charuco(crop, annots)
|
|
k2d = annots['keypoints2d']
|
|
for i in range(self.nJoints):
|
|
param['annots']['keypoints2d'][i][0] = k2d[i][0] + start[0]
|
|
param['annots']['keypoints2d'][i][1] = k2d[i][1] + start[1]
|
|
param['annots']['keypoints2d'][i][2] = k2d[i][2]
|
|
|
|
def detect_charuco(self, crop, annots):
|
|
cfg = {
|
|
'long': 6,
|
|
'short': 4,
|
|
'squareLength': 0.128,
|
|
'aruco_len': 0.1,
|
|
'aruco_type': '4X4_50'
|
|
}
|
|
from easymocap.annotator.chessboard import CharucoBoard
|
|
board = CharucoBoard(**cfg)
|
|
board.detect(crop, annots)
|
|
|
|
def annot_example(path, sub, args):
|
|
# define datasets
|
|
calib = Matcher(path, mode=args.mode, args=args)
|
|
dataset = ImageFolder(path, image=args.image, sub=sub,
|
|
annot=args.annot, no_annot=False, ext=args.ext,
|
|
share_annot=True,
|
|
max_per_folder=-1)
|
|
# define visualize
|
|
vis_funcs = [vis_point, vis_line, vis_bbox, calib.vis, resize_to_screen, plot_text]
|
|
key_funcs = {
|
|
' ': calib.add,
|
|
'b': calib.back,
|
|
'z': calib.add_conf,
|
|
'p': calib.add_point_by_2lines,
|
|
'c': calib.clear_point,
|
|
'C': calib.clear,
|
|
'x': calib.clear,
|
|
'e': calib.detect,
|
|
}
|
|
# construct annotations
|
|
annotator = AnnotBase(
|
|
dataset=dataset,
|
|
key_funcs=key_funcs,
|
|
vis_funcs=vis_funcs)
|
|
while annotator.isOpen:
|
|
annotator.run()
|
|
|
|
if __name__ == "__main__":
|
|
from easymocap.annotator import load_parser, parse_parser
|
|
parser = load_parser()
|
|
parser.add_argument('--pattern', type=lambda x: (int(x.split(',')[0]), int(x.split(',')[1])),
|
|
help='The pattern of the chessboard', default=(9, 6))
|
|
parser.add_argument('--mode', type=str, default='chessboard')
|
|
args = parse_parser(parser)
|
|
mywarn('[annot_calib] The default patter is {}'.format(args.pattern))
|
|
if len(args.sub) == 0:
|
|
args.sub = ['']
|
|
for sub in args.sub:
|
|
annot_example(args.path, sub, args) |