update detect chessboard

apps/calibration/detect_chessboard.py

@@ -2,72 +2,105 @@
   @ Date: 2021-07-16 20:13:57
   @ Author: Qing Shuai
   @ LastEditors: Qing Shuai
-  @ LastEditTime: 2021-07-21 19:56:38
+  @ LastEditTime: 2022-05-11 20:41:10
-  @ FilePath: /EasyMocap/apps/calibration/detect_chessboard.py
+  @ 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 = imgname.replace(ext, '.json')
-        annname = join(path, annname)
+        annname = join(path, 'chessboard', annname)
-        if os.path.exists(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):
+def _detect_chessboard(datas, path, image, out, pattern):
-    create_chessboard(path, pattern, gridSize, ext=args.ext)
+    for imgname, annotname in datas:
-    dataset = ImageFolder(path, annot='chessboard', ext=args.ext)
+        # imgname, annotname = dataset[i]
-    dataset.isTmp = False
-    if args.silent:
-        trange = range(len(dataset))
-    else:
-        trange = tqdm(range(len(dataset)))
-    for i in trange:
-        imgname, annotname = dataset[i]
         # 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:
+            mywarn('[Info] Cannot find chessboard in {}'.format(imgname))
-                print('Cannot find {}'.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)
+        if isinstance(show, np.ndarray):
             cv2.imwrite(outname, show)
 
-def detect_chessboard_sequence(path, out, pattern, gridSize, args):
+def detect_chessboard(path, image, out, pattern, gridSize, args):
-    create_chessboard(path, pattern, gridSize, ext=args.ext)
+    create_chessboard(path, image, pattern, gridSize, ext=args.ext, overwrite=args.overwrite3d)
-    subs = sorted(os.listdir(join(path, 'images')))
+    dataset = ImageFolder(path, image=image, annot='chessboard', ext=args.ext)
-    for sub in subs:
+    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=np.bool)
+    found = np.zeros(nFrames, dtype=bool)
-        visited = np.zeros(nFrames, dtype=np.bool)
+    visited = np.zeros(nFrames, dtype=bool)
     proposals = []
     init_step = args.max_step
     min_step = args.min_step
@@ -76,7 +109,8 @@ def detect_chessboard_sequence(path, out, pattern, gridSize, args):
             proposals.append([nf, nf+init_step])
     while len(proposals) > 0:
         left, right = proposals.pop(0)
-            print('Check [{}, {}]'.format(left, right))
+        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
@@ -84,18 +118,24 @@ def detect_chessboard_sequence(path, out, pattern, gridSize, args):
                 # 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('Cannot find {}'.format(imgname))
+                        print('[Info] Cannot find chessboard in {}'.format(imgname))
                     found[nf] = False
                     continue
                 found[nf] = True
-                    outname = join(out, imgname.replace(path + '/images/', ''))
+                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:
@@ -105,10 +145,50 @@ def detect_chessboard_sequence(path, out, pattern, gridSize, args):
         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
+        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)