add __init__

apps/calibration/detect_chessboard.py

@@ -1,3 +1,10 @@
+'''
+  @ Date: 2021-07-16 20:13:57
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-07-17 19:25:00
+  @ FilePath: /EasyMocapRelease/apps/calibration/detect_chessboard.py
+'''
 # detect the corner of chessboard
 from easymocap.annotator.file_utils import getFileList, read_json, save_json
 from tqdm import tqdm
@@ -29,11 +36,15 @@ def create_chessboard(path, pattern, gridSize, ext):
         else:
             save_json(annname, template)
 
-def detect_chessboard(path, out, pattern, gridSize):
+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
-    for i in tqdm(range(len(dataset))):
+    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)
@@ -41,22 +52,72 @@ def detect_chessboard(path, out, pattern, gridSize):
         show = findChessboardCorners(img, annots, pattern)
         save_json(annotname, annots)
         if show is None:
+            if args.debug:
+                print('Cannot find {}'.format(imgname))
             continue
         outname = join(out, imgname.replace(path + '/images/', ''))
         os.makedirs(os.path.dirname(outname), exist_ok=True)
         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')))
+    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 = 50
+        min_step = 1
+        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 > min_step:
+                proposals.append((left, mid))
+            if right - mid > min_step:
+                proposals.append((mid, right))
+
 if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument('path', type=str)
-    parser.add_argument('--out', type=str)
+    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])),
         help='The pattern of the chessboard', default=(9, 6))
     parser.add_argument('--grid', type=float, default=0.1, 
         help='The length of the grid size (unit: meter)')
+    parser.add_argument('--silent', action='store_true')
     parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--seq', action='store_true')
     args = parser.parse_args()
-
+    if args.seq:
-    detect_chessboard(args.path, args.out, pattern=args.pattern, gridSize=args.grid)
+        detect_chessboard_sequence(args.path, 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)

easymocap/__init__.py

@@ -0,0 +1,7 @@
+'''
+  @ Date: 2021-07-17 19:24:14
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-07-17 19:24:14
+  @ FilePath: /EasyMocapRelease/easymocap/__init__.py
+'''

easymocap/annotator/basic_dataset.py

@@ -1,32 +1,92 @@
+'''
+  @ Date: 2021-04-15 16:57:53
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-07-14 22:15:26
+  @ FilePath: /EasyMocap/easymocap/annotator/basic_dataset.py
+'''
 from os.path import join
+import os
+import shutil
 from .file_utils import getFileList
 
 class ImageFolder:
-    def __init__(self, path, sub=None, annot='annots') -> None:
+    def __init__(self, path, sub=None, image='images', annot='annots', no_annot=False, ext='.jpg', remove_tmp=True) -> None:
         self.root = path
+        self.image = image
+        self.annot = annot
+        self.image_root = join(path, self.image)
+        self.annot_root = join(path, self.annot)
+        if not os.path.exists(self.annot_root):
+            no_annot = True
+        self.annot_root_tmp = join(path, self.annot + '_tmp')
+        if os.path.exists(self.annot_root_tmp) and remove_tmp:
+            shutil.rmtree(self.annot_root_tmp)
+        if sub is None:
+            self.imgnames = getFileList(self.image_root, ext=ext)
+            if not no_annot:
+                self.annnames = getFileList(self.annot_root, ext='.json')
+        else:
+            self.imgnames = getFileList(join(self.image_root, sub), ext=ext)
+            self.imgnames = [join(sub, name) for name in self.imgnames]
+            if not no_annot:
+                self.annnames = getFileList(join(self.annot_root, sub), ext='.json')
+                self.annnames = [join(sub, name) for name in self.annnames]
+                length = min(len(self.imgnames), len(self.annnames))
+                self.imgnames = self.imgnames[:length]
+                self.annnames = self.annnames[:length]
+                # assert len(self.imgnames) == len(self.annnames)
+        self.isTmp = True
+        self.no_annot = no_annot
+    
+    def __getitem__(self, index):
+        imgname = join(self.image_root, self.imgnames[index])
+        if self.no_annot:
+            annname = None
+        else:
+            if self.isTmp:
+                annname = join(self.annot_root_tmp, self.annnames[index])
+            else:
+                annname = join(self.annot_root, self.annnames[index])
+        return imgname, annname
+    
+    def __len__(self):
+        return len(self.imgnames)
+    
+    def __str__(self) -> str:
+        return '{}: {} images'.format(self.root, len(self))
+
+class MVBase:
+    def __init__(self, path, subs, annot='annots') -> None:
+        self.root = path
+        self.subs = subs
         self.image = 'images'
         self.annot = annot
         self.image_root = join(path, self.image)
         self.annot_root = join(path, self.annot)
         self.annot_root_tmp = join(path, self.annot + '_tmp')
-        if sub is None:
+        assert len(subs) > 0, subs
-            self.imgnames = getFileList(self.image_root, ext='.jpg')
+        self.imgnames, self.annnames = {}, {}
-            self.annnames = getFileList(self.annot_root, ext='.json')
+        for sub in subs:
-        else:
+            imgnames = getFileList(join(self.image_root, sub), ext='.jpg')
-            self.imgnames = getFileList(join(self.image_root, sub), ext='.jpg')
+            annnames = getFileList(join(self.annot_root, sub), ext='.json')
-            self.annnames = getFileList(join(self.annot_root, sub), ext='.json')
+            imgnames = [join(sub, name) for name in imgnames]
-            self.imgnames = [join(sub, name) for name in self.imgnames]
+            annnames = [join(sub, name) for name in annnames]
-            self.annnames = [join(sub, name) for name in self.annnames]
+            self.imgnames[sub] = imgnames
+            self.annnames[sub] = annnames
         self.isTmp = True
-        assert len(self.imgnames) == len(self.annnames)
+        
     
     def __getitem__(self, index):
-        imgname = join(self.image_root, self.imgnames[index])
+        imgnames, annnames = {}, {}
-        if self.isTmp:
+        for sub in self.subs:
-            annname = join(self.annot_root_tmp, self.annnames[index])
+            imgnames[sub] = join(self.image_root, self.imgnames[sub][index])
-        else:
+            if self.isTmp:
-            annname = join(self.annot_root, self.annnames[index])
+                annname = join(self.annot_root_tmp, self.annnames[sub][index])
-        return imgname, annname
+            else:
+                annname = join(self.annot_root, self.annnames[sub][index])
+            annnames[sub] = annname
+        return imgnames, annnames
     
     def __len__(self):
-        return len(self.imgnames)
+        return len(self.imgnames[self.subs[0]])