EasyMocap/apps/annotation/annot_clip.py

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