diff --git a/apps/calibration/detect_chessboard.py b/apps/calibration/detect_chessboard.py index ce2633a..057f6e6 100644 --- a/apps/calibration/detect_chessboard.py +++ b/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() - - detect_chessboard(args.path, args.out, pattern=args.pattern, gridSize=args.grid) \ No newline at end of file + if args.seq: + 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) \ No newline at end of file diff --git a/easymocap/__init__.py b/easymocap/__init__.py new file mode 100644 index 0000000..b1eeae1 --- /dev/null +++ b/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 +''' diff --git a/easymocap/annotator/basic_dataset.py b/easymocap/annotator/basic_dataset.py index 3bec9f3..f22313a 100644 --- a/easymocap/annotator/basic_dataset.py +++ b/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: - self.imgnames = getFileList(self.image_root, ext='.jpg') - self.annnames = getFileList(self.annot_root, ext='.json') - else: - self.imgnames = getFileList(join(self.image_root, sub), ext='.jpg') - self.annnames = getFileList(join(self.annot_root, sub), ext='.json') - self.imgnames = [join(sub, name) for name in self.imgnames] - self.annnames = [join(sub, name) for name in self.annnames] + assert len(subs) > 0, subs + self.imgnames, self.annnames = {}, {} + for sub in subs: + imgnames = getFileList(join(self.image_root, sub), ext='.jpg') + annnames = getFileList(join(self.annot_root, sub), ext='.json') + imgnames = [join(sub, name) for name in imgnames] + annnames = [join(sub, name) for name in 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]) - if self.isTmp: - annname = join(self.annot_root_tmp, self.annnames[index]) - else: - annname = join(self.annot_root, self.annnames[index]) - return imgname, annname + imgnames, annnames = {}, {} + for sub in self.subs: + imgnames[sub] = join(self.image_root, self.imgnames[sub][index]) + if self.isTmp: + annname = join(self.annot_root_tmp, self.annnames[sub][index]) + else: + annname = join(self.annot_root, self.annnames[sub][index]) + annnames[sub] = annname + return imgnames, annnames def __len__(self): - return len(self.imgnames) \ No newline at end of file + return len(self.imgnames[self.subs[0]]) \ No newline at end of file