Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
import os
2
import sys
3
import cv2
4
import argparse
5
from insightface_func.face_detect_crop_single import Face_detect_crop
6
from pathlib import Path
7
from tqdm import tqdm
8

9
def main(args):
10
    app = Face_detect_crop(name='antelope', root='./insightface_func/models')
11
    app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
12
    crop_size = 224
13

14
    dirs = os.listdir(args.path_to_dataset)
15
    for i in tqdm(range(len(dirs))):
16
        d = os.path.join(args.path_to_dataset, dirs[i])
17
        dir_to_save = os.path.join(args.save_path, dirs[i])
18
        Path(dir_to_save).mkdir(parents=True, exist_ok=True)
19
        
20
        image_names = os.listdir(d)
21
        for image_name in image_names:
22
            try:
23
                image_path = os.path.join(d, image_name)
24
                image = cv2.imread(image_path)
25
                cropped_image, _ = app.get(image, crop_size)
26
                cv2.imwrite(os.path.join(dir_to_save, image_name), cropped_image[0])
27
            except:
28
                pass
29
        
30
    
31
if __name__ == "__main__":
32
    parser = argparse.ArgumentParser()
33
    
34
    parser.add_argument('--path_to_dataset', default='./VggFace2/VGG-Face2/data/preprocess_train', type=str)
35
    parser.add_argument('--save_path', default='./VggFace2-crop', type=str)
36
    
37
    args = parser.parse_args()
38
    
39
    main(args)
40

41