update preprocess

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
+  @ LastEditTime: 2022-07-12 13:39:25
-  @ FilePath: /EasyMocap/apps/preprocess/extract_image.py
+  @ 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:
+        if args.scale != 1 and args.transpose != -1:
-            other_cmd += '-vf transpose={}'.format(args.transpose)
+            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()
 

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