EasyMocap/apps/annotation/annot_calib.py

'''
  @ 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)