facenet代码实现

来源:互联网 发布:淘宝店铺招牌图片童装 编辑:程序博客网 时间:2024/06/05 09:19

详细代码见百度云

train_classifier.py解读

训练一个facenet人脸分类器,调用已有的模型训练自己的模型


import tensorflow as tf

import numpy as np
import cv2
import os
from os.path import join as pjoin
import sys
import copy
import detect_face
import nn4 as network
import matplotlib.pyplot as plt
import sklearn
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn import metrics  
from sklearn.externals import joblib

#face detection parameters
minsize = 20
threshold = [ 0.6, 0.7, 0.7 ] 
factor = 0.709 
model_dir='./model_check_point/model-20160506.ckpt-500000'  
image_size=96 
pool_type='MAX' 
use_lrn=False 
seed=42
batch_size= None 

def to_rgb(img):
    w, h = img.shape
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, 0] = ret[:, :, 1] = ret[:, :, 2] = img
    return ret

def read_img(person_dir,f):
    img=cv2.imread(pjoin(person_dir, f))    
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)        
    if gray.ndim == 2:
        img = to_rgb(gray)
    return img

def load_data(data_dir):
    data = {}
    pics_ctr = 0
    for guy in os.listdir(data_dir):
        person_dir = pjoin(data_dir, guy)      
        curr_pics = [read_img(person_dir, f) for f in os.listdir(person_dir)]                
        data[guy] = curr_pics        
    return data

# KNN Classifier  
def knn_classifier(train_x, train_y):  
    from sklearn.neighbors import KNeighborsClassifier  
    model = KNeighborsClassifier()  
    model.fit(train_x, train_y)  
    return model 

if __name__ == '__main__':
    #建立人脸检测模型,加载参数
    print('Creating networks and loading parameters')
    gpu_memory_fraction=1.0
    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
        with sess.as_default():
            pnet, rnet, onet = detect_face.create_mtcnn(sess, './model_check_point/')

   
    #建立facenet embedding模型
    print('建立facenet embedding模型')
    tf.Graph().as_default()
    sess = tf.Session()
    images_placeholder = tf.placeholder(tf.float32, shape=(batch_size,image_size,image_size, 3), name='input')
    phase_train_placeholder = tf.placeholder(tf.bool, name='phase_train')

    embeddings = network.inference(images_placeholder, pool_type,use_lrn,1.0,phase_train=

phase_train_placeholder)

   

    ema = tf.train.ExponentialMovingAverage(1.0)

    saver = tf.train.Saver(ema.variables_to_restore())
    model_checkpoint_path='./model_check_point/model-20160506.ckpt-500000'
    saver.restore(sess, model_checkpoint_path)
    print('facenet embedding模型建立完毕')


    ###### train_dir containing one subdirectory per image class 
    #should like this:
    #-->train_dir:
    #     --->pic_female:
    #            f1.jpg
    #            f2.jpg
    #            ...
    #     --->pic_male:
    #           m1.jpg
    #           m2.jpg
    #            ...

    

    data_dir='./train/'

    data=load_data(data_dir)
    keys=[]
    for key in data.iterkeys():
        keys.append(key)
        print('foler:{},image numbers:{}'.format(key,len(data[key])))    
    train_x=[]
    train_y=[]

    for x in data[keys[0]]:

        #人脸检测,调用已定义的detect_face函数

        bounding_boxes, _ = detect_face.detect_face(x, minsize, pnet, rnet, onet, threshold, factor)
        nrof_faces = bounding_boxes.shape[0]#number of faces 
        for face_position in bounding_boxes:
            face_position=face_position.astype(int)
            #print(face_position[0:4])
            cv2.rectangle(x, (face_position[0], face_position[1]), (face_position[2], face_position[3]), (0, 255, 0), 2)
            crop=x[face_position[1]:face_position[3],face_position[0]:face_position[2],]  
            crop = cv2.resize(crop, (96, 96), interpolation=cv2.INTER_CUBIC )

            #print(crop.shape)     
            crop_data=crop.reshape(-1,96,96,3)
            #print(crop_data.shape)
            emb_data = sess.run([embeddings],feed_dict={images_placeholder: np.array(crop_data), phase_train_placeholder: False })[0]               
            train_x.append(emb_data)
            train_y.append(0)
    print(len(train_x))
    
    for y in data[keys[1]]:
        bounding_boxes, _ = detect_face.detect_face(y, minsize, pnet, rnet, onet, threshold, factor)
        nrof_faces = bounding_boxes.shape[0]    
        for face_position in bounding_boxes:
            face_position=face_position.astype(int)
            #print(face_position[0:4])
            #cv2.rectangle(y, (face_position[0], face_position[1]), (face_position[2], face_position[3]), (0, 255, 0), 2)
            crop=y[face_position[1]:face_position[3],face_position[0]:face_position[2],]
            crop = cv2.resize(crop, (96, 96), interpolation=cv2.INTER_CUBIC )
            crop_data=crop.reshape(-1,96,96,3)
            #print(crop_data.shape)
        
            emb_data = sess.run([embeddings], feed_dict={images_placeholder: np.array(crop_data), phase_train_placeholder: False })[0]
              
            train_x.append(emb_data)
            train_y.append(1)
    print(len(train_x))
    print('搞完了,样本数为:{}'.format(len(train_x)))


    #train/test split
    train_x=np.array(train_x)
    train_x=train_x.reshape(-1,128)
    train_y=np.array(train_y)
    print(train_x.shape)
    print(train_y.shape)

    #从训练集中分离出训练集与测试集
    X_train, X_test, y_train, y_test = train_test_split(train_x, train_y, test_size=.3, random_state=42)
    print(X_train.shape,y_train.shape,X_test.shape,y_test.shape)  

    #knn分类器

    classifiers = knn_classifier 

    model = classifiers(X_train,y_train)  
    predict = model.predict(X_test)  

    accuracy = metrics.accuracy_score(y_test, predict)  
    print ('accuracy: %.2f%%' % (100 * accuracy)  )    
    #save model
    joblib.dump(model, './model_check_point/knn_classifier_gender.model')
    


        

test_classifier.py解读

调用训练好的模型进行测试


import tensorflow as tf
from tensorflow.core.protobuf import saver_pb2
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
from os.path import join as pjoin
import sys
import copy
import detect_face
import nn4 as network
import random
import sklearn
from sklearn.externals import joblib


#face detection parameters
minsize = 20 
threshold = [ 0.6, 0.7, 0.7 ]  
factor = 0.709 
model_dir='./model_check_point/model.ckpt-500000')
image_size=96 
pool_type='MAX' 
use_lrn=False 
seed=42,# "Random seed."
batch_size= None 
frame_interval=1 


def to_rgb(img):
  w, h = img.shape
  ret = np.empty((w, h, 3), dtype=np.uint8)
  ret[:, :, 0] = ret[:, :, 1] = ret[:, :, 2] = img
  return ret


if __name__ == '__main__':
    #restore mtcnn model
    print('Creating networks and loading parameters')
    gpu_memory_fraction=1.0
    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
        with sess.as_default():
            pnet, rnet, onet = detect_face.create_mtcnn(sess, './model_check_point/')

    #restore facenet model
    print('建立facenet embedding模型')
    tf.Graph().as_default()
    sess = tf.Session()
    images_placeholder = tf.placeholder(tf.float32, shape=(batch_size, image_size, image_size, 3), name='input')
    phase_train_placeholder = tf.placeholder(tf.bool, name='phase_train')

            # 创建网络图:除了全连接层和损失层
    embeddings = network.inference(images_placeholder, pool_type, use_lrn, 1.0, phase_train=

phase_train_placeholder)


    ema = tf.train.ExponentialMovingAverage(1.0)
    saver = tf.train.Saver(ema.variables_to_restore())
    model_checkpoint_path='./model_check_point/model-20160506.ckpt-500000'
    saver.restore(sess, model_checkpoint_path)
    print('facenet embedding模型建立完毕')


    #restore pre-trained knn classifier
    model = joblib.load('./model_check_point/knn_classifier_gender.model')
    print('knn classifier loaded 建立完毕')
    image = cv2.imread(sys.argv[1])
    find_results=[]
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
    if gray.ndim == 2:
        img = to_rgb(gray)
    bounding_boxes,points = detect_face.detect_face(img, minsize,pnet, rnet, onet, threshold, factor) 
    nrof_faces = bounding_boxes.shape[0]#number of faces
      
    for face_position in bounding_boxes:
        face_position=face_position.astype(int)                     
        cv2.rectangle(image, (face_position[0], face_position[1]),(face_position[2], face_position[3]), (0, 255, 0), 2)
            
        crop=img[face_position[1]:face_position[3],face_position[0]:face_position[2],]  
        crop = cv2.resize(crop, (96, 96), interpolation=cv2.INTER_CUBIC )
        
        data=crop.reshape(-1,96,96,3)     
        emb_data = sess.run([embeddings], feed_dict={images_placeholder: np.array(data), phase_train_placeholder: False })[0]
            
        predict = model.predict(emb_data) 
        print predict              
        if predict[0]==0:
            result = 'female'
        elif predict[0]==1:
            result = 'male'


        #cv2.putText在图片上添加水印
        cv2.putText(image,result, (face_position[0]-10,face_position[1]-10), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (255, 0 ,0),thickness = 2, lineType = 2)


    # show result
    image = cv2.resize(image, (640, 480), interpolation=cv2.INTER_CUBIC)
    cv2.imshow("Show Result",image)
    cv2.waitKey(0)


  
            
       

    

原创粉丝点击