yolov5人脸识别(yolov5-facenet-svm)

yolov5出来的这段时间,挺火,自己试着跑了以下,速度精度确实不错,相比yolov3性能要高,相比yolov4这个咱不敢说。以下用yolov5做一个人脸识别的demo。
本文章项目代码github地址https://github.com/BlackFeatherQQ/FaceRecognition
一、yolov5
yolov5的原理这个就不介绍了,网上一查一大堆,github官方源码地址https://github.com/ultralytics/yolov5
现在来谈谈yolov5在这个人来能识别demo中的作用(框出人脸):
怎么才能框出人脸呢,yolov5官方预训练权重只能框出整个人,要想框出人脸只能在其基础之上再训练,训练走起:
1)数据从哪里来:
开源人脸数据celeba与wideface,本人选择使用的是celeba,原因:简单、易于训练、人脸识别项目精度已经够用了
2)怎么训练:
具体请看yolov5训练详解
3)获取人脸的框

二、facenet
FaceNet亮点:
  1. 利用DNN直接学习到从原始图片到欧氏距离空间的映射,从而使得在欧式空间里的距离的度量直接关联着人脸相似度;
  2. 引入triplet损失函数,使得模型的学习能力更高效。
具体原理在此不再介绍,百度可以找到很多,我们在此介绍它的应用。
下载facenet预训练权重https://github.com/davidsandberg/facenet
yolov5人脸识别(yolov5-facenet-svm)_第1张图片
代码不需要下载,太繁杂,如果你想要训练的话可以下载,此处我们只使用facenet.py已经够用了。
三、SVM
支持向量机(support vector machines, SVM)是一种二分类模型,它的基本模型是定义在特征空间上的间隔最大的线性分类器,间隔最大使它有别于感知机;SVM还包括核技巧,这使它成为实质上的非线性分类器。SVM的的学习策略就是间隔最大化,可形式化为一个求解凸二次规划的问题,也等价于正则化的合页损失函数的最小化问题。SVM的的学习算法就是求解凸二次规划的最优化算法。
需要我们去实现吗?当然不需要,sklearn内置,直接调用即可
四、数据制作
(1)人脸数据
yolov5人脸识别(yolov5-facenet-svm)_第2张图片
此处做了一点数据增强(镜像、高斯模糊等)
yolov5人脸识别(yolov5-facenet-svm)_第3张图片
每个文件夹的名称对应每个人的名字(注意文件名不要使用中文(opencv暂不支持中文编码,编译器等可能也会报错误或者出现乱码问题))
(2)人脸特征提取保存

 picture_path = "D:/code_data/face_recognition/images" #用户图片数据
 model_path = "D:/code_data/facenet/20180402-114759"  #facenet模型权重
 database_path = "D:/code_data/face_recognition/npz/Database.npz"  #打包成npz文件
 face2database(picture_path,model_path,database_path) #step1
def face2database(picture_path, model_path, database_path, batch_size=90, image_size=160):
    # 提取特征到数据库
    # picture_path为人脸文件夹的所在路径
    # model_path为facenet模型路径
    # database_path为人脸数据库路径
    with tf.Graph().as_default():
        with tf.Session() as sess:
            dataset = facenet.get_dataset(picture_path)
            paths, labels = facenet.get_image_paths_and_labels(dataset)
            print('Number of classes: %d' % len(dataset))
            print('Number of images: %d' % len(paths))
            # Load the model
            print('Loading feature extraction model')
            facenet.load_model(model_path)
            # Get input and output tensors
            images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
            embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
            phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
            embedding_size = embeddings.get_shape()[1]
            # Run forward pass to calculate embeddings
            print('Calculating features for images')
            nrof_images = len(paths)
            nrof_batches_per_epoch = int(math.ceil(1.0 * nrof_images / batch_size))
            emb_array = np.zeros((nrof_images, embedding_size))
            for i in range(nrof_batches_per_epoch):
                start_index = i * batch_size
                end_index = min((i + 1) * batch_size, nrof_images)
                paths_batch = paths[start_index:end_index]
                images = facenet.load_data(paths_batch, False, False, image_size)
                feed_dict = {images_placeholder: images, phase_train_placeholder: False}
                emb_array[start_index:end_index, :] = sess.run(embeddings, feed_dict=feed_dict)
            np.savez(database_path, emb=emb_array, lab=labels)
            print("数据库特征提取完毕!")
            # emb_array里存放的是图片特征,labels为对应的标签

(3)训练SVM分类器

picture_path = "D:/code_data/face_recognition/images" #用户图片数据
model_path = "D:/code_data/facenet/20180402-114759"  #facenet模型权重
SVCpath = "D:/code_data/face_recognition/pkl/SVCmodel.pkl"  #通过SVM训练保存到pkl
ClassifyTrainSVC(database_path,SVCpath) #step2
def ClassifyTrainSVC(database_path, SVCpath):
    # database_path为人脸数据库
    # SVCpath为分类器储存的位置
    Database = np.load(database_path)
    name_lables = Database['lab']
    embeddings = Database['emb']
    name_unique = np.unique(name_lables)
    labels = []
    for i in range(len(name_lables)):
        for j in range(len(name_unique)):
            if name_lables[i] == name_unique[j]:
                labels.append(j)
    print('Training classifier')
    model = SVC(kernel='linear', probability=True)
    model.fit(embeddings, labels)
    with open(SVCpath, 'wb') as outfile:
        pickle.dump((model, name_unique), outfile)
        print('Saved classifier model to file "%s"' % SVCpath)

五、识别
整理一下思路:
数据库中的人脸数据处理:通过fanet提取特征保存在Database.npz中,利用特征训练SVM分类器保存在SVCmodel.pkl中。
yolov5通过训练celeba数据集可以框出图像中的人脸,我们把框出的人脸拿出来丢到facenet中会生成人脸特征,把人脸特征和我们数据库中的进行对比(SVM或者直接计算特征距离)算出相似度。
把yolov5与facenet代码关联起来,如下面的代码(支持图片、文件夹、视频,自动判别文件类型)

def setOPT():
    # 文件配置
    # *******************************************************

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='D:/py/FaceRecognition/weights/last.pt', help='model.pt path')
    parser.add_argument('--source', type=str, default='C:/Users/lieweiai/Desktop/video/26321934-1-192.mp4', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='../inference/output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
    parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--facenet-model-path',type=str,default='D:/code_data/facenet/20180402-114759',help='miss facenet-model')
    parser.add_argument('--svc-path', type=str, default='D:/code_data/face_recognition/pkl/SVCmodel.pkl',
                        help='miss svc')
    parser.add_argument('--database-path', type=str, default='D:/code_data/face_recognition/npz/Database.npz',
                        help='miss database')
    opt = parser.parse_args()
    opt.img_size = check_img_size(opt.img_size)
    print(opt)

    return opt
def detect(opt, save_img=False):
    out, source, weights, view_img, save_txt, imgsz, facenet_model_path, svc_path, database_path = \
        opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.facenet_model_path, opt.svc_path, opt.database_path

    webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')

    # Initialize
    device = torch_utils.select_device(opt.device)
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    google_utils.attempt_download(weights)
    model = torch.load(weights, map_location=device)['model'].float()  # load to FP32
    # torch.save(torch.load(weights, map_location=device), weights)  # update model if SourceChangeWarning
    # model.fuse()
    model.to(device).eval()
    if half:
        model.half()  # to FP16

    # Second-stage classifier
    classify = False
    if classify:
        modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
        modelc.to(device).eval()

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = True
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz)
    else:
        save_img = True
        dataset = LoadImages(source, img_size=imgsz)

    # Get names and colors
    names = model.names if hasattr(model, 'names') else model.modules.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    # Run inference
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once

    # ************************************************************************

    with tf.Graph().as_default():
        with tf.Session() as sess:
            # Load the model
            print('Loading feature extraction model')
            facenet.load_model(facenet_model_path)
            with open(svc_path, 'rb') as infile:
                (classifymodel, class_names) = pickle.load(infile)
            print('Loaded classifier model from file "%s"' % svc_path)

            # Get input and output tensors
            images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
            embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
            phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
            embedding_size = embeddings.get_shape()[1]
            Database = np.load(database_path)

            corpbbox = None


        # ************************************************************

            for path, img, im0s, vid_cap in dataset:
                img = torch.from_numpy(img).to(device)
                img = img.half() if half else img.float()  # uint8 to fp16/32
                img /= 255.0  # 0 - 255 to 0.0 - 1.0
                if img.ndimension() == 3:
                    img = img.unsqueeze(0)

                # Inference
                t1 = torch_utils.time_synchronized()
                pred = model(img, augment=opt.augment)[0]

                # Apply NMS
                pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
                t2 = torch_utils.time_synchronized()

                # Apply Classifier
                if classify:
                    pred = apply_classifier(pred, modelc, img, im0s)

                # Process detections
                for i, det in enumerate(pred):  # detections per image
                    if webcam:  # batch_size >= 1
                        p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
                    else:
                        p, s, im0 = path, '', im0s


                    # ******************************
                    image = Image.fromarray(cv2.cvtColor(im0, cv2.COLOR_BGR2RGB))
                    image = np.array(image)
                    img_size = np.array(image.shape)[0:2]
                    # ********************************

                    save_path = str(Path(out) / Path(p).name)
                    s += '%gx%g ' % img.shape[2:]  # print string
                    gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  #  normalization gain whwh
                    if det is not None and len(det):
                        # Rescale boxes from img_size to im0 size
                        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                        # Print results
                        for c in det[:, -1].unique():
                            n = (det[:, -1] == c).sum()  # detections per class
                            s += '%g %ss, ' % (n, names[int(c)])  # add to string

                        # Write results
                        for *xyxy, conf, cls in det:
                            if save_txt:  # Write to file
                                xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(
                                    -1).tolist()  # normalized xywh
                                with open(save_path[:save_path.rfind('.')] + '.txt', 'a') as file:
                                    file.write(('%g ' * 5 + '\n') % (cls, *xywh))  # label format

                            if save_img or view_img:  # Add bbox to image

                                # ***************************************************

                                x1 = np.maximum(int(xyxy[0]) - 16, 0)
                                y1 = np.maximum(int(xyxy[1]) - 16, 0)
                                x2 = np.minimum(int(xyxy[2]) + 16, img_size[1])
                                y2 = np.minimum(int(xyxy[3]) + 16, img_size[0])
                                crop_img = image[y1:y2, x1:x2]
                                scaled = np.array(Image.fromarray(crop_img).resize((160, 160)))
                                # scaled = misc.imresize(crop_img, (160, 160), interp='bilinear')
                                img = load_image(scaled, False, False, 160)
                                img = np.reshape(img, (-1, 160, 160, 3))
                                feed_dict = {images_placeholder: img, phase_train_placeholder: False}
                                embvecor = sess.run(embeddings, feed_dict=feed_dict)
                                embvecor = np.array(embvecor)

                                # 利用SVM对人脸特征进行分类
                                predictions = classifymodel.predict_proba(embvecor)
                                best_class_indices = np.argmax(predictions, axis=1)
                                tmp_lable = class_names[best_class_indices]
                                best_class_probabilities = predictions[
                                    np.arange(len(best_class_indices)), best_class_indices]
                                print(class_names, predictions)

                                if best_class_probabilities < 0.95:
                                    tmp_lable = "others"

                                print(tmp_lable)
                                # ***************************************************

                                # label = '%s %.2f' % (names[int(cls)], conf)
                                label = '%s %.2f' % (tmp_lable, best_class_probabilities)
                                plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

                    # Print time (inference + NMS)
                    print('%sDone. (%.3fs)' % (s, t2 - t1))
                    cv2.imshow("",im0)
                    cv2.waitKey(5)
                # ***********************************************************



                # ************************************************************************



                    # Stream results
                    if view_img:
                        cv2.imshow(p, im0)
                        if cv2.waitKey(1) == ord('q'):  # q to quit
                            raise StopIteration

                    # Save results (image with detections)
                    if save_img:
                        if dataset.mode == 'images':
                            cv2.imwrite(save_path, im0)
                        else:
                            if vid_path != save_path:  # new video
                                vid_path = save_path
                                if isinstance(vid_writer, cv2.VideoWriter):
                                    vid_writer.release()  # release previous video writer

                                fps = vid_cap.get(cv2.CAP_PROP_FPS)
                                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                                vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
                            vid_writer.write(im0)

            if save_txt or save_img:
                print('Results saved to %s' % os.getcwd() + os.sep + out)
                if platform == 'darwin':  # MacOS
                    os.system('open ' + save_path)

            print('Done. (%.3fs)' % (time.time() - t0))




def face2database(picture_path, model_path, database_path, batch_size=90, image_size=160):
    # 提取特征到数据库
    # picture_path为人脸文件夹的所在路径
    # model_path为facenet模型路径
    # database_path为人脸数据库路径
    with tf.Graph().as_default():
        with tf.Session() as sess:
            dataset = facenet.get_dataset(picture_path)
            paths, labels = facenet.get_image_paths_and_labels(dataset)
            print('Number of classes: %d' % len(dataset))
            print('Number of images: %d' % len(paths))
            # Load the model
            print('Loading feature extraction model')
            facenet.load_model(model_path)
            # Get input and output tensors
            images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
            embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
            phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
            embedding_size = embeddings.get_shape()[1]
            # Run forward pass to calculate embeddings
            print('Calculating features for images')
            nrof_images = len(paths)
            nrof_batches_per_epoch = int(math.ceil(1.0 * nrof_images / batch_size))
            emb_array = np.zeros((nrof_images, embedding_size))
            for i in range(nrof_batches_per_epoch):
                start_index = i * batch_size
                end_index = min((i + 1) * batch_size, nrof_images)
                paths_batch = paths[start_index:end_index]
                images = facenet.load_data(paths_batch, False, False, image_size)
                feed_dict = {images_placeholder: images, phase_train_placeholder: False}
                emb_array[start_index:end_index, :] = sess.run(embeddings, feed_dict=feed_dict)
            np.savez(database_path, emb=emb_array, lab=labels)
            print("数据库特征提取完毕!")
            # emb_array里存放的是图片特征,labels为对应的标签


def ClassifyTrainSVC(database_path, SVCpath):
    # database_path为人脸数据库
    # SVCpath为分类器储存的位置
    Database = np.load(database_path)
    name_lables = Database['lab']
    embeddings = Database['emb']
    name_unique = np.unique(name_lables)
    labels = []
    for i in range(len(name_lables)):
        for j in range(len(name_unique)):
            if name_lables[i] == name_unique[j]:
                labels.append(j)
    print('Training classifier')
    model = SVC(kernel='linear', probability=True)
    model.fit(embeddings, labels)
    with open(SVCpath, 'wb') as outfile:
        pickle.dump((model, name_unique), outfile)
        print('Saved classifier model to file "%s"' % SVCpath)


# 图片预处理阶段
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 prewhiten(x):
    mean = np.mean(x)
    std = np.std(x)
    std_adj = np.maximum(std, 1.0 / np.sqrt(x.size))
    y = np.multiply(np.subtract(x, mean), 1 / std_adj)
    return y


def crop(image, random_crop, image_size):
    if image.shape[1] > image_size:
        sz1 = int(image.shape[1] // 2)
        sz2 = int(image_size // 2)
        if random_crop:
            diff = sz1 - sz2
            (h, v) = (np.random.randint(-diff, diff + 1), np.random.randint(-diff, diff + 1))
        else:
            (h, v) = (0, 0)
        image = image[(sz1 - sz2 + v):(sz1 + sz2 + v), (sz1 - sz2 + h):(sz1 + sz2 + h), :]
    return image


def flip(image, random_flip):
    if random_flip and np.random.choice([True, False]):
        image = np.fliplr(image)
    return image


def load_image(image_old, do_random_crop, do_random_flip, image_size, do_prewhiten=True):
    if image_old.ndim == 2:
        image_old = to_rgb(image_old)
    if do_prewhiten:
        image_old = prewhiten(image_old)
    image_old = crop(image_old, do_random_crop, image_size)
    image_old = flip(image_old, do_random_flip)
    return image_old

【提示】github上的代码不包含yolov5训练那一块,recognition/test.py是用来做人脸识别的入口,在里面你可以更改相应配置。
我的数据库只做了胡歌、刘亦菲、刘德华的数据,未在数据库中的会显示为others,附几张效果图吧
yolov5人脸识别(yolov5-facenet-svm)_第4张图片
yolov5人脸识别(yolov5-facenet-svm)_第5张图片

你可能感兴趣的:(人脸识别,python,深度学习,人工智能,机器学习,pytorch)