diff --git a/apps/preprocess/extract_image.py b/apps/preprocess/extract_image.py index b3d3809..9fcbd0d 100644 --- a/apps/preprocess/extract_image.py +++ b/apps/preprocess/extract_image.py @@ -2,8 +2,8 @@ @ Date: 2021-08-19 22:06:13 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-10-23 16:02:43 - @ FilePath: /EasyMocap/apps/preprocess/extract_image.py + @ LastEditTime: 2022-07-12 13:39:25 + @ FilePath: /EasyMocapPublic/apps/preprocess/extract_image.py ''' # extract image from videos import os @@ -25,14 +25,22 @@ def extract_images(path, ffmpeg, image): sub = '.'.join(os.path.basename(videoname).split('.')[:-1]) sub = sub.replace(args.strip, '') outpath = join(path, image, sub) + if os.path.exists(outpath) and (len(os.listdir(outpath)) > 10 or len(os.listdir(outpath)) == args.num) and not args.restart: + continue os.makedirs(outpath, exist_ok=True) other_cmd = '' if args.num != -1: other_cmd += '-vframes {}'.format(args.num) - if args.transpose != -1: - other_cmd += '-vf transpose={}'.format(args.transpose) + if args.scale != 1 and args.transpose != -1: + other_cmd += ' -vf "transpose={transpose},scale=iw/{scale}:ih/{scale}"'.format(scale=args.scale,transpose=args.transpose) + elif args.scale != 1: + other_cmd += ' -vf "scale=iw/{scale}:ih/{scale}'.format(scale=args.scale) + elif args.transpose != -1: + other_cmd += ' -vf transpose={}'.format(args.transpose) cmd = '{} -i {} {} -q:v 1 -start_number 0 {}/%06d.jpg'.format( ffmpeg, videoname, other_cmd, outpath) + if not args.debug: + cmd += ' -loglevel quiet' run(cmd) if __name__ == "__main__": @@ -42,8 +50,10 @@ if __name__ == "__main__": parser.add_argument('--strip', type=str, default='') parser.add_argument('--image', type=str, default='images') parser.add_argument('--num', type=int, default=-1) + parser.add_argument('--scale', type=int, default=1) parser.add_argument('--transpose', type=int, default=-1) parser.add_argument('--ffmpeg', type=str, default='ffmpeg') + parser.add_argument('--restart', action='store_true') parser.add_argument('--debug', action='store_true') args = parser.parse_args() diff --git a/apps/preprocess/extract_track.py b/apps/preprocess/extract_track.py new file mode 100644 index 0000000..eea370f --- /dev/null +++ b/apps/preprocess/extract_track.py @@ -0,0 +1,247 @@ +from easymocap.annotator.file_utils import read_json, save_json +from easymocap.config import load_object_from_cmd +import numpy as np +from easymocap.mytools.debug_utils import log, myerror, mywarn, run_cmd +from tqdm import tqdm +import os +from os.path import join + +class Tracker: + def __init__(self, missing_frame=10, thres_iou=0.5) -> None: + self.cache = {} + self.max_id = -1 + self.time = 0 + self.dist_mode = 'bbox' + self.min_accept_dist = thres_iou + self.missing_frame = missing_frame + self.failed = {} + + def step(self): + self.time += 1 + removelist = [] + for pid, track in self.cache.items(): + if self.time - track['end_time'] > self.missing_frame: + mywarn('[{:06d}] Delete person {:3d} with {:6d} frames'.format(self.time, pid, track['end_time'] - track['start_time'])) + removelist.append(pid) + for pid in removelist: + self.failed[pid] = self.cache.pop(pid) + + def init(self, data): + pid = data['personID'] + self.max_id = max(self.max_id, pid) + self.cache[pid] = { + 'start_time': self.time, + 'end_time': self.time, + 'missing_frame': [], + 'infos': [data] + } + return True, pid + + def update(self, data, pid): + track = self.cache[pid] + if self.time == track['end_time'] + 1 or self.time != 1: + track['end_time'] = self.time + else: + mywarn('[{:06d}] Refind person {:3d} from {:06d}'.format(self.time, pid, track['end_time'])) + track['end_time'] = self.time + for f in range(track['end_time'] + 1, self.time): + track['missing_frame'].append(f) + track['infos'].append(data) + + def calculate_distance(self, data, infos): + # TODO: require the last frame + info = infos[-1] + if self.dist_mode == 'bbox': + bbox_now = data['bbox'] + bbox_pre = info['bbox'] + area_now = (bbox_now[2] - bbox_now[0])*(bbox_now[3]-bbox_now[1]) + area_pre = (bbox_pre[2] - bbox_pre[0])*(bbox_pre[3]-bbox_pre[1]) + # compute IOU + # max of left + xx1 = max(bbox_now[0], bbox_pre[0]) + yy1 = max(bbox_now[1], bbox_pre[1]) + # min of right + xx2 = min(bbox_now[0+2], bbox_pre[0+2]) + yy2 = min(bbox_now[1+2], bbox_pre[1+2]) + # w h + w = max(0, xx2 - xx1) + h = max(0, yy2 - yy1) + over = (w*h)/(area_pre+area_now-w*h) + distance = 1 - over + return distance + + def track(self, data): + keys = list(self.cache.keys()) + distance = np.zeros(len(keys)) + 999. + for ikey, key in enumerate(keys): + if self.cache[key]['end_time'] == self.time: + # already assigned in current frame + continue + else: + dist = self.calculate_distance(data, self.cache[key]['infos']) + distance[ikey] = dist + if (distance > 10).all(): + # all tracks have been assigned in current frame + return False, -1 + best_id = distance.argmin() + if distance[best_id] > self.min_accept_dist: + mywarn('[{:06d}] Tracking failed with distance {}'.format(self.time, distance[best_id])) + return False, -1 + else: + self.cache[keys[best_id]] + self.update(data, keys[best_id]) + return True, keys[best_id] + + def add(self, data): + flag, pid = self.track(data) + if not flag: + log('[{:06d}] Create person {:3d}'.format(self.time, self.max_id+1)) + data['personID'] = self.max_id + 1 + flag, pid = self.init(data) + return flag, pid + + def report(self): + removelist = [] + for pid, track in self.cache.items(): + if track['end_time'] - track['start_time'] < self.missing_frame: + removelist.append(pid) + for pid in removelist: + self.failed[pid] = self.cache.pop(pid) + # success + log('- Tracked detection:') + for pid, track in self.cache.items(): + log('{:3d} [{:6d}->{:6d}], missing {}'.format(pid, track['start_time'], track['end_time'], track['missing_frame'])) + +def track2d(datas): + # sort the first frame by size + annots0 = datas['annots'][0] + len_first_frame = len(annots0) + tracker = Tracker(thres_iou=args.thres_iou) + annots0.sort(key=lambda x:-(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1])) + for i, annot in enumerate(annots0): + annot['personID'] = i + for nf, annots in enumerate(datas['annots']): + if nf == 0: + # new the tracker + for annot in annots: + flag, pid = tracker.init(annot) + continue + # track all the frames + tracker.step() + annots0.sort(key=lambda x:-(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1])) + for annot in annots: + # greedy match + flag, pid = tracker.add(annot) + if flag: + annot['personID'] = pid + from easymocap.annotator.file_utils import save_annot + nFrames = len(data['annname']) + for nf in tqdm(range(nFrames), desc='writing track'): + annname = data['annname'][nf] + annots = data['annots'][nf] + annots.sort(key=lambda x:x['personID']) + annots_origin = read_json(annname) + annots_origin['annots'] = annots + save_annot(annname, annots_origin) + tracker.report() + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser(description= + ''' + For common usage: + python3 apps/preprocess/extract_track.py ${data} + For fast motion: + python3 apps/preprocess/extract_track.py ${data} --thres_iou 0.8 + ''') + parser.add_argument('path', type=str) + parser.add_argument('--out', type=str) + parser.add_argument('--subs', type=str, default=[], nargs='+') + parser.add_argument('--max', type=int, default=-1) + parser.add_argument('--thres_iou', type=float, default=0.5) + parser.add_argument('--annot_track', action='store_true') + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + opt_data = ['args.path', args.path] + annotbase = join(args.path, 'database.json') + if not os.path.exists(annotbase): + save_json(annotbase, {}) + if False: + trackname = join(args.path, 'track.json') + track_info = read_json(trackname) + annot_info = {} + for sub, flag in track_info.items(): + if not os.path.exists(join(args.path, 'images', sub)): + continue + annot_info[sub] = { + 'skip': 0, + 'tracked': flag, + 'detected': flag, + 'reconstructed': 0, + } + save_json(annotbase, annot_info) + import ipdb;ipdb.set_trace() + exit() + annotbase = join(args.path, 'database.json') + track_info = read_json(annotbase) + + if not os.path.exists(annotbase): + save_json(annotbase, {}) + + if len(args.subs) >0: + opt_data.append('args.subs') + opt_data.append(args.subs) + else: + # check subs + subs_all = sorted(os.listdir(join(args.path, 'annots'))) + subs = [] + for sub in subs_all: + # check if the sub has been tracked + if sub not in track_info: + track_info[sub] = { + 'skip': 0, + 'tracked': 0, + 'detected': 0, + 'reconstructed': 0, + } + if track_info[sub]['skip'] or track_info[sub]['tracked']: + mywarn('- skip {}'.format(sub)) + continue + len_img = len(os.listdir(join(args.path, 'images', sub))) + len_ann = len(os.listdir(join(args.path, 'annots', sub))) + if len_img != len_ann: + mywarn('- skip {} as no enough detections'.format(sub)) + continue + subs.append(sub) + opt_data.append('args.subs') + opt_data.append(subs) + dataset = load_object_from_cmd('config/data/multivideo-mp.yml', opt_data) + for idx in range(len(dataset)): + # try: + data = dataset[idx] + # except: + # myerror('- Failed to load {}'.format(dataset.subs[idx])) + # continue + track2d(data) + imgname = data['imgname'][0] + sub = os.path.basename(os.path.dirname(imgname)) + valid = input('Does this track right? [y/n]') + if valid == 'y': + track_info[sub]['tracked'] = 1 + track_info[sub]['detected'] = 1 + elif args.annot_track: + # detect of reclip this + cmd = f'python3 apps/annotation/annot_track.py {args.path} --sub {sub}' + run_cmd(cmd) + save_json(annotbase, track_info) + track_info = read_json(annotbase) + # check track + failed_subs = [] + for sub, flag in track_info.items(): + if not flag['tracked']: + failed_subs.append(sub) + if len(failed_subs) > 0: + mywarn('- Success subs: {}'.format(len(track_info.keys()) - len(failed_subs))) + mywarn('- Failed subs: {}'.format(failed_subs)) + log('Run the following command to annotate the failed tracks:') + log('python3 apps/annotation/annot_track.py ${{data}} --sub {}'.format(' '.join(failed_subs))) \ No newline at end of file