217 lines
7.4 KiB
Python
217 lines
7.4 KiB
Python
'''
|
|
@ Date: 2021-06-09 09:57:23
|
|
@ Author: Qing Shuai
|
|
@ LastEditors: Qing Shuai
|
|
@ LastEditTime: 2022-07-14 21:37:34
|
|
@ FilePath: /EasyMocapPublic/apps/annotation/annot_clip.py
|
|
'''
|
|
# 功能：
|
|
# 1. 快速预览图像
|
|
# 2. 设置起点
|
|
# 3. 设置终点
|
|
# 不兼容的接口：没有标注文件
|
|
from easymocap.mytools.debug_utils import myerror, mywarn, run_cmd
|
|
from easymocap.mytools.vis_base import plot_line
|
|
from easymocap.annotator.basic_annotator import AnnotBase, parse_parser
|
|
from easymocap.annotator import ImageFolder
|
|
from easymocap.annotator import plot_text
|
|
from easymocap.annotator.basic_visualize import capture_screen, resize_to_screen
|
|
from easymocap.mytools import read_json, save_json
|
|
from easymocap.annotator.basic_keyboard import get_any_move
|
|
from os.path import join
|
|
import os
|
|
import numpy as np
|
|
import cv2
|
|
|
|
class Clips:
|
|
def __init__(self, path) -> None:
|
|
self.temp = join(path, 'clips.json')
|
|
if os.path.exists(self.temp):
|
|
self.annots = read_json(self.temp)
|
|
else:
|
|
self.annots = {}
|
|
self.start_ = None
|
|
self.end_ = None
|
|
self.clips = []
|
|
self.sub_ = None
|
|
|
|
@property
|
|
def sub(self):
|
|
return self.sub_
|
|
|
|
@sub.setter
|
|
def sub(self, value):
|
|
self.sub_ = value
|
|
if value in self.annots.keys():
|
|
self.clips = self.annots[value]
|
|
else:
|
|
self.annots[value] = []
|
|
self.clips = self.annots[value]
|
|
self.print(0)
|
|
|
|
def start(self, annotator, **kwargs):
|
|
self.start_ = annotator.frame
|
|
print('>>> Start clip from frame {:6d}'.format(annotator.frame))
|
|
|
|
def end(self, annotator, **kwargs):
|
|
self.end_ = annotator.frame
|
|
print('>>> End clip from frame {:6d}'.format(annotator.frame))
|
|
|
|
def add(self, annotator, **kwargs):
|
|
if self.start_ is None:
|
|
print('[clip] Please check the start!')
|
|
return 0
|
|
if self.end_ is None:
|
|
print('[clip] Please check the end!')
|
|
return 0
|
|
print('[{}, {})'.format(self.start_, self.end_))
|
|
self.clips.append([self.start_, self.end_])
|
|
self.start_ = None
|
|
self.end_ = None
|
|
|
|
def delete(self, annotator, **kwargs):
|
|
frame = annotator.frame
|
|
ind = -1
|
|
for i, (start, end) in enumerate(self.clips):
|
|
if frame > start and frame < end:
|
|
ind = i
|
|
break
|
|
else:
|
|
print('[clip] current not in any clip')
|
|
return 0
|
|
self.clips.pop(ind)
|
|
|
|
def print(self, annotator, **kwargs):
|
|
print('{}: '.format(self.sub))
|
|
for (start, end) in self.clips:
|
|
print(' - [{}, {})'.format(start, end))
|
|
|
|
def save(self):
|
|
save_json(self.temp, self.annots)
|
|
|
|
def vis_clips(self, img, frame, nFrames, **kwargs):
|
|
COL_CLIP = (0, 0, 255)
|
|
COL_NEW = (0, 0, 255)
|
|
width = img.shape[1]
|
|
pos = lambda x: int(width*(x+1)/nFrames)
|
|
lw = 12
|
|
# 可视化标注的clips
|
|
for (start, end) in self.clips:
|
|
plot_line(img, (pos(start), lw/2), (pos(end), lw/2), lw, COL_CLIP)
|
|
# 可视化当前的标注
|
|
if self.start_ is not None:
|
|
top = pos(self.start_)
|
|
pts = np.array([[top, lw], [top-lw, lw*4], [top, lw*4]])
|
|
cv2.fillPoly(img, [pts], COL_NEW)
|
|
if self.end_ is not None:
|
|
top = pos(self.end_)
|
|
pts = np.array([[top, lw], [top, lw*4], [top+lw, lw*4]])
|
|
cv2.fillPoly(img, [pts], COL_NEW)
|
|
return img
|
|
|
|
def annot_example(path, sub, skip=False):
|
|
# define datasets
|
|
# define visualize
|
|
if not os.path.exists(join(path, 'images', sub)):
|
|
mywarn('[annot] No such sub: {}'.format(sub))
|
|
return 0
|
|
clip = Clips(path)
|
|
vis_funcs = [resize_to_screen, plot_text, clip.vis_clips, capture_screen]
|
|
clip.sub = sub
|
|
if skip and len(clip.clips) > 0:
|
|
return 0
|
|
key_funcs = {
|
|
'j': clip.start,
|
|
'k': clip.end,
|
|
'l': clip.add,
|
|
'x': clip.delete,
|
|
'v': clip.print,
|
|
'w': get_any_move(-10),
|
|
's': get_any_move(10),
|
|
'f': get_any_move(100),
|
|
'g': get_any_move(-100)
|
|
}
|
|
|
|
dataset = ImageFolder(path, sub=sub, no_annot=True)
|
|
print('[Info] Totally {} frames'.format(len(dataset)))
|
|
# construct annotations
|
|
annotator = AnnotBase(
|
|
dataset=dataset,
|
|
key_funcs=key_funcs,
|
|
vis_funcs=vis_funcs)
|
|
while annotator.isOpen:
|
|
annotator.run()
|
|
clip.save()
|
|
|
|
def copy_clips(path, out):
|
|
from tqdm import tqdm
|
|
import shutil
|
|
from easymocap.mytools.debug_utils import log, mywarn, mkdir
|
|
temp = join(path, 'clips.json')
|
|
assert os.path.exists(temp), temp
|
|
annots = read_json(temp)
|
|
for key, clips in tqdm(annots.items()):
|
|
for start, end in clips:
|
|
outname = '{}+{:06d}+{:06d}'.format(key, start, end)
|
|
outdir = join(out, 'images', outname)
|
|
if os.path.exists(outdir) and len(os.listdir(outdir)) == end - start:
|
|
mywarn('[copy] Skip {}'.format(outname))
|
|
continue
|
|
# check the input image
|
|
srcname0 = join(path, 'images', key, '{:06d}.jpg'.format(start))
|
|
srcname1 = join(path, 'images', key, '{:06d}.jpg'.format(end))
|
|
if not os.path.exists(srcname0) or not os.path.exists(srcname1):
|
|
myerror('[copy] No such file: {}, {}'.format(srcname0, srcname1))
|
|
log('[copy] {}'.format(outname))
|
|
mkdir(outdir)
|
|
# copy the images
|
|
for nnf, nf in enumerate(tqdm(range(start, end), desc='copy {}'.format(outname))):
|
|
srcname = join(path, 'images', key, '{:06d}.jpg'.format(nf))
|
|
dstname = join(outdir, '{:06d}.jpg'.format(nnf))
|
|
shutil.copyfile(srcname, dstname)
|
|
|
|
def copy_mv_clips(path, out):
|
|
temp = join(path, 'clips.json')
|
|
assert os.path.exists(temp), temp
|
|
annots = read_json(temp)
|
|
clips = list(annots.values())[0]
|
|
for start, end in clips:
|
|
if out is None:
|
|
outdir = path + '+{:06d}+{:06d}'.format(start, end)
|
|
else:
|
|
outdir = out + '+{:06d}+{:06d}'.format(start, end)
|
|
print(outdir)
|
|
cmd = f'python3 scripts/preprocess/copy_dataset.py {path} {outdir} --start {start} --end {end}'
|
|
if len(args.sub) > 0:
|
|
cmd += ' --subs {}'.format(' '.join(args.sub))
|
|
if args.strip is not None:
|
|
cmd += ' --strip {}'.format(args.strip)
|
|
run_cmd(cmd)
|
|
|
|
if __name__ == "__main__":
|
|
from easymocap.annotator import load_parser, parse_parser
|
|
parser = load_parser()
|
|
parser.add_argument('--strip', type=str, default=None)
|
|
parser.add_argument('--copy', action='store_true')
|
|
parser.add_argument('--skip', action='store_true')
|
|
parser.add_argument('--mv', action='store_true')
|
|
parser.add_argument('--sub_ignore', type=str, nargs='+', default=[])
|
|
args = parse_parser(parser)
|
|
|
|
args.sub = [i for i in args.sub if i not in args.sub_ignore]
|
|
|
|
if args.copy:
|
|
print(args.path, args.out)
|
|
if args.mv:
|
|
copy_mv_clips(args.path, args.out)
|
|
else:
|
|
if args.out is None:
|
|
myerror('[copy] No output path')
|
|
exit(0)
|
|
copy_clips(args.path, args.out)
|
|
else:
|
|
if args.mv:
|
|
annot_example(args.path, sub=args.sub[0], skip=args.skip)
|
|
else:
|
|
for sub in args.sub:
|
|
annot_example(args.path, sub=sub, skip=args.skip) |