MTCNN批量生成人脸图像

直接上代码了

#coding:utf-8
import sys
sys.path.append('..')
from Detection.MtcnnDetector import MtcnnDetector
from Detection.detector import Detector
from Detection.fcn_detector import FcnDetector
from train_models.mtcnn_model import P_Net, R_Net, O_Net
from prepare_data.loader import TestLoader
import cv2
import os
import numpy as np
test_mode = "ONet"
thresh = [0.9, 0.6, 0.7]
min_face_size = 24
stride = 2
slide_window = False
shuffle = False
detectors = [None, None, None]
prefix = ['../data/MTCNN_model/PNet_landmark/PNet', '../data/MTCNN_model/RNet_landmark/RNet', '../data/MTCNN_model/ONet_landmark/ONet']
epoch = [18, 14, 16]
batch_size = [2048, 256, 16]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
# load pnet model
if slide_window:
    PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
    PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet

# load rnet model
if test_mode in ["RNet", "ONet"]:
    RNet = Detector(R_Net, 24, batch_size[1], model_path[1])
    detectors[1] = RNet

# load onet model
if test_mode == "ONet":
    ONet = Detector(O_Net, 48, batch_size[2], model_path[2])
    detectors[2] = ONet

mtcnn_detector = MtcnnDetector(detectors=detectors, min_face_size=min_face_size,
                               stride=stride, threshold=thresh, slide_window=slide_window)
gt_imdb = []
#gt_imdb.append("35_Basketball_Basketball_35_515.jpg")
#imdb_ = dict()"
#imdb_['image'] = im_path
#imdb_['label'] = 5
#path = "lala"
path = "Z:/Attack_20181112_P"
PathSave = "F:/Attack_20181112_P"
StartN = 0
savesize = (96+2*8)

FindFN = 0#folder number
PathFind = []
import fnmatch
def ReadSaveAddr2(Stra):
    path = Stra
    for dirpath,dirnames,filenames in os.walk(path):
        #for filename in filenames:
        a_list = fnmatch.filter(os.listdir(dirpath),'*.jpg')
        if len(a_list):
            PathFind.append(dirpath)
            continue
        a_list = fnmatch.filter(os.listdir(dirpath),'*.bmp')
        if len(a_list):
            PathFind.append(dirpath)
            continue
        a_list = fnmatch.filter(os.listdir(dirpath),'*.png')
        if len(a_list):
            PathFind.append(dirpath)
            continue
            
ReadSaveAddr2(path)

for i in range(len(PathFind)):
    gt_imdb = []
    for item in os.listdir(PathFind[i]):
        imagepath = os.path.join(PathFind[i],item)
        image = cv2.imread(imagepath)
        if image is None:
            continue
        else:
            gt_imdb.append(os.path.join(PathFind[i],item))
    test_data = TestLoader(gt_imdb)
    all_boxes,landmarks = mtcnn_detector.detect_face(test_data)
    count = 0
    for imagepath in gt_imdb:
        print (imagepath)
        image = cv2.imread(imagepath)
        for bbox in all_boxes[count]:
            x0 = bbox[0]
            y0 = bbox[1]
            x1 = bbox[2]
            y1 = bbox[3]
            w = x1-x0
            h = y1-y0
            w = w/12
            h = h/12
            x0 = x0-w
            y0 = y0-h
            x1 = x1+w
            y1 = y1+h
            bbox[0] = x0
            bbox[1] = y0
            bbox[2] = x1
            bbox[3] = y1
            if int(bbox[0])<0:
                continue
            if int(bbox[1])<0:
                continue
            size = image.shape
            if int(bbox[2])>=size[1]:
                continue
            if int(bbox[3])>=size[0]:
                continue
            roi = image[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2]) ]
            roi = cv2.resize(roi,(savesize,savesize))
            n = str(StartN)
            s = n.zfill(6)
            SaveN = PathSave+'/'+s+'.jpg'
            StartN = StartN+1
            cv2.imwrite(SaveN,roi)
            break
        count = count + 1