Spaces:
Running on Zero
Running on Zero
| import argparse | |
| import cv2 | |
| import glob | |
| import numpy as np | |
| import os | |
| import torch | |
| from tqdm import tqdm | |
| from facexlib.detection import init_detection_model | |
| from facexlib.tracking.sort import SORT | |
| def main(args): | |
| detect_interval = args.detect_interval | |
| margin = args.margin | |
| face_score_threshold = args.face_score_threshold | |
| save_frame = True | |
| if save_frame: | |
| colors = np.random.rand(32, 3) | |
| # init detection model and tracker | |
| det_net = init_detection_model('retinaface_resnet50', half=False) | |
| tracker = SORT(max_age=1, min_hits=2, iou_threshold=0.2) | |
| print('Start track...') | |
| # track over all frames | |
| frame_paths = sorted(glob.glob(os.path.join(args.input_folder, '*.jpg'))) | |
| pbar = tqdm(total=len(frame_paths), unit='frames', desc='Extract') | |
| for idx, path in enumerate(frame_paths): | |
| img_basename = os.path.basename(path) | |
| frame = cv2.imread(path) | |
| img_size = frame.shape[0:2] | |
| # detection face bboxes | |
| with torch.no_grad(): | |
| bboxes = det_net.detect_faces(frame, 0.97) | |
| additional_attr = [] | |
| face_list = [] | |
| for idx_bb, bbox in enumerate(bboxes): | |
| score = bbox[4] | |
| if score > face_score_threshold: | |
| bbox = bbox[0:5] | |
| det = bbox[0:4] | |
| # face rectangle | |
| det[0] = np.maximum(det[0] - margin, 0) | |
| det[1] = np.maximum(det[1] - margin, 0) | |
| det[2] = np.minimum(det[2] + margin, img_size[1]) | |
| det[3] = np.minimum(det[3] + margin, img_size[0]) | |
| face_list.append(bbox) | |
| additional_attr.append([score]) | |
| trackers = tracker.update(np.array(face_list), img_size, additional_attr, detect_interval) | |
| pbar.update(1) | |
| pbar.set_description(f'{idx}: detect {len(bboxes)} faces in {img_basename}') | |
| # save frame | |
| if save_frame: | |
| for d in trackers: | |
| d = d.astype(np.int32) | |
| cv2.rectangle(frame, (d[0], d[1]), (d[2], d[3]), colors[d[4] % 32, :] * 255, 3) | |
| if len(face_list) != 0: | |
| cv2.putText(frame, 'ID : %d DETECT' % (d[4]), (d[0] - 10, d[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.75, colors[d[4] % 32, :] * 255, 2) | |
| cv2.putText(frame, 'DETECTOR', (5, 45), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (1, 1, 1), 2) | |
| else: | |
| cv2.putText(frame, 'ID : %d' % (d[4]), (d[0] - 10, d[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, | |
| colors[d[4] % 32, :] * 255, 2) | |
| save_path = os.path.join(args.save_folder, img_basename) | |
| cv2.imwrite(save_path, frame) | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--input_folder', help='Path to the input folder', type=str) | |
| parser.add_argument('--save_folder', help='Path to save visualized frames', type=str, default=None) | |
| parser.add_argument( | |
| '--detect_interval', | |
| help=('how many frames to make a detection, trade-off ' | |
| 'between performance and fluency'), | |
| type=int, | |
| default=1) | |
| # if the face is big in your video ,you can set it bigger for easy tracking | |
| parser.add_argument('--margin', help='add margin for face', type=int, default=20) | |
| parser.add_argument( | |
| '--face_score_threshold', help='The threshold of the extracted faces,range 0 < x <=1', type=float, default=0.85) | |
| args = parser.parse_args() | |
| os.makedirs(args.save_folder, exist_ok=True) | |
| main(args) | |
| # add verification | |
| # remove last few frames | |