fulyvv/EasyMocap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add __init__

Browse Source
This commit is contained in:
shuaiqing 2021-07-17 19:27:51 +08:00
parent 401ce08d50
commit 2a3edecda4
3 changed files with 150 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
apps/calibration/detect_chessboard.py
View File

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

7
easymocap/__init__.py Normal file
View File

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

72
easymocap/annotator/basic_dataset.py
View File

@ -1,27 +1,49 @@
'''
@ 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 = 'images'
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='.jpg')
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='.jpg')
self.annnames = getFileList(join(self.annot_root, sub), ext='.json')
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
assert len(self.imgnames) == len(self.annnames)
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:
@ -30,3 +52,41 @@ class ImageFolder:
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')
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
def __getitem__(self, index):
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[self.subs[0]])