图片分类

# coding: utf-8
import sys
import os
import dlib
import glob
import cv2
import tkinter

# root = tkinter.Tk()
# root.minsize(400, 400)
# menubar = tkinter.Menu(root)
# filemenu = tkinter.Menu(menubar, tearoff=0)
# menubar.add_cascade(label='命令', menu=filemenu)
# menubar.add_cascade(label='工具')
# menubar.add_cascade(label='收藏夹')
# menubar.add_cascade(label='选项')
# menubar.add_cascade(label='帮助')
# root.config(menu=menubar)
def index():
    # 指定路径
    current_path = os.getcwd()
    model_path = current_path + '\\model\\'
    shape_predictor_model = model_path + 'shape_predictor_5_face_landmarks.dat'
    face_rec_model = model_path + 'dlib_face_recognition_resnet_model_v1.dat'
    face_folder = current_path + '\\faces\\'
    output_folder = current_path + '\\output\\'

    # 导入模型
    detector = dlib.get_frontal_face_detector()
    shape_detector = dlib.shape_predictor(shape_predictor_model)
    face_recognizer = dlib.face_recognition_model_v1(face_rec_model)

    # 为后面操作方便，建了几个列表
    descriptors = []
    images = []
    # 遍历faces文件夹中所有的图片
    for f in glob.glob(os.path.join(face_folder, "*.jpg")):
        print('处理文件：{}'.format(f))
        # 读取图片
        img = cv2.imread(f)
        # 转换到rgb颜色空间
        img2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 检测人脸
        dets = detector(img2, 1)
        print("检测到的人脸数: {}".format(len(dets)))

        # 遍历所有的人脸
        for index, face in enumerate(dets):
            # 检测人脸特征点
            shape = shape_detector(img2, face)
            # 投影到128D
            face_descriptor = face_recognizer.compute_face_descriptor(img2, shape)

            # 保存相关信息
            descriptors.append(face_descriptor)
            images.append((img2, shape))

    # 聚类
    labels = dlib.chinese_whispers_clustering(descriptors, 0.5)
    #labels
    print("标签: {}".format(labels))
    num_classes = len(set(labels))
    #Number of clusters
    print("分类数: {}".format(num_classes))

    # 为了方便操作，用字典类型保存
    face_dict = {}
    for i in range(num_classes):
        face_dict[i] = []
    #print (face_dict)
    for i in range(len(labels)):
        face_dict[labels[i]].append(images[i])

    #print (face_dict.keys())
    # 遍历字典，保存结果
    for key in face_dict.keys():
        file_dir = os.path.join(output_folder, str(key))
        if not os.path.isdir(file_dir):
            os.makedirs(file_dir)

        for index, (image, shape) in enumerate(face_dict[key]):
            file_path = os.path.join(file_dir, 'face_' + str(index))
            print("图片位置：",file_path)
            dlib.save_face_chip(image, shape, file_path)
    
# button1 =tkinter.Button(text="图片分类",command=index)
# button1.place(x=250,y=150,relheight=0.1)
# label1 =tkinter.Label()
# label1.place()
#
# root.mainloop()

if __name__ == '__main__':
#def main():
    while True:
        print("")
        print("="*40)
        print("1.faces文件夹是存放分类图片的位置")
        print("2.output文件夹是整理后存放图片的位置")
        print("3.退出请输入：quit")
        print("="*40)
        print("")
        print("是否进行分类？")
        a = input("请输入Y进行确认： ")
        if a == "Y" or a == "y":
            index()
        elif a =="quit":
            break
        else:
            print("请重新输入")