python-learning-notes/select_images.py at master · PanJinquan/python-learning-notes · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
# -*-coding: utf-8 -*-
"""
    @Project: python-learning-notes
    @File   : select_images.py
    @Author : panjq
    @E-mail : pan_jinquan@163.com
    @Date   : 2019-09-23 16:04:20
"""

import os
import cv2
import numpy as np
import shutil
import os.path
from utils import file_processing, image_processing
from libs.ultra_ligh_face.ultra_ligh_face import UltraLightFaceDetector


def ramdom_select_image_dir(image_dir, dest_dir):
    select_nums = 100
    image_id = file_processing.get_sub_directory_list(image_dir)
    for id in image_id:
        image_list = file_processing.get_files_list(os.path.join(image_dir, id),
                                                    postfix=['*.jpg', "*.jpeg", '*.png', "*.JPG"])
        image_list = np.random.permutation(image_list)[:select_nums]
        for src_path in image_list:
            basename = os.path.basename(src_path)
            dest_path = file_processing.create_dir(dest_dir, id, basename)
            shutil.copy(src_path, dest_path)


def select_image_dir(image_dir, dest_dir):
    image_id = file_processing.get_sub_directory_list(image_dir)
    for id in image_id:
        image_list = file_processing.get_files_list(os.path.join(image_dir, id),
                                                    postfix=['*.jpg', "*.jpeg", '*.png', "*.bmp"])
        for src_path in image_list:
            basename = os.path.basename(src_path)
            index = basename.split(".")[0].split("_")[1]
            if index == "0":
                dest_path = file_processing.create_dir(dest_dir, id, basename)
                # shutil.copy(src_path, dest_path)
                file_processing.move_file(src_path, dest_path)


def select_facebank_detect(image_dir, dest_dir, id_nums=None, detect_face=True):
    if detect_face:
        # model_path = "../../face_detection/face_detection_rbf.pth"
        # model_path = "/media/dm/dm1/git/python-learning-notes/libs/ultra_ligh_face/face_detection_rbf.pth"
        model_path = "/home/panjinquan/project/python-learning-notes//libs/ultra_ligh_face/face_detection_rbf.pth"
        network = "RFB"
        confidence_threshold = 0.85
        nms_threshold = 0.3
        top_k = 500
        keep_top_k = 750
        device = "cuda:0"
        detector = UltraLightFaceDetector(model_path=model_path,
                                          network=network,
                                          confidence_threshold=confidence_threshold,
                                          nms_threshold=nms_threshold,
                                          top_k=top_k,
                                          keep_top_k=keep_top_k,
                                          device=device)
    per_nums = 1
    image_id = file_processing.get_sub_directory_list(image_dir)
    nums_images = len(image_id)
    print("have ID:{}".format(nums_images))
    if id_nums:
        id_nums = min(id_nums, nums_images)
        image_id = image_id[:id_nums]
    print("select ID:{}".format(len(image_id)))

    for id in image_id:
        image_list = file_processing.get_files_list(os.path.join(image_dir, id),
                                                    postfix=['*.jpg', "*.jpeg", '*.png', "*.JPG"])
        count = 0
        for src_path in image_list:
            basename = os.path.basename(src_path)
            if detect_face:
                bgr_image = cv2.imread(src_path)
                bboxes, scores, landms = detector.detect(bgr_image, isshow=True)
                if not len(bboxes) == 1:
                    print("no face:{}".format(src_path))
                    continue
            if count >= per_nums:
                break
            count += 1
            dest_path = file_processing.create_dir(dest_dir, id, basename)
            file_processing.copy_file(src_path, dest_path)


def select_facebank(image_dir, dest_dir, id_nums=10):
    per_nums = 1
    image_id = file_processing.get_sub_directory_list(image_dir)
    nums_images = len(image_id)
    print("have ID:{}".format(nums_images))
    if id_nums:
        id_nums = min(id_nums, nums_images)
        image_id = image_id[:id_nums]
    print("select ID:{}".format(len(image_id)))
    for id in image_id:
        image_list = file_processing.get_files_list(os.path.join(image_dir, id),
                                                    postfix=['*.jpg', "*.jpeg", '*.png', "*.JPG"])
        count = 0
        for src_path in image_list:
            basename = os.path.basename(src_path)
            if count >= per_nums:
                break
            count += 1
            dest_path = file_processing.create_dir(dest_dir, id, basename)
            file_processing.copy_file(src_path, dest_path)


def image_to_facebank(image_dir, dest_dir):
    from xpinyin import Pinyin
    p = Pinyin()
    image_list = file_processing.get_files_list(image_dir,
                                                postfix=['*.jpg', "*.jpeg", '*.png', "*.JPG"])
    nums_images = len(image_list)
    print("have ID:{}".format(nums_images))
    for image_path in image_list:
        basename = os.path.basename(image_path)
        id_name = basename.split(".")[0]
        id_name = p.get_pinyin(id_name, '')
        dest_path = file_processing.create_dir(dest_dir, id_name, basename)
        file_processing.copy_file(image_path, dest_path)


if __name__ == "__main__":
    # image_dir = '/media/dm/dm1/FaceDataset/lexue/lexue1/val-src'
    # dest_dir = '/media/dm/dm1/FaceDataset/lexue/lexue1/val'
    # ramdom_select_image_dir(image_dir, dest_dir)

    # image_dir = '/media/dm/dm1/FaceDataset/X4/CASIA-FaceV5/trainval'
    # dest_dir = '/media/dm/dm1/FaceDataset/X4/CASIA-FaceV5/facebank'
    # select_image_dir(image_dir, dest_dir)

    image_dir = "/data0/panjinquan/FaceData/DMFR_V1"
    dest_dir = "/data0/panjinquan/FaceData/facebank_DMFR_V1"
    select_facebank_detect(image_dir, dest_dir, id_nums=None, detect_face=True)