mtcnn 终章

import numpy as np
import torch
from PIL import Image
from nms import nms
from PROnat import Pnet,RNet,ONet
from utils import *
from .first_stage import run_first_stage
device = torch.device("cuda:0" if torch.cuda.is_available else "cpu")
class MTCNN():

    def __int__(self):
        self.pnet = Pnet().to(device)
        self.rnet = RNet().to(device)
        self.onet = ONet().to(device)
        self.pnet.eval()
        self.rnet.eval()
        self.onet.eval()

    def detect_faces(self,image,min_face_size = 20.0
                     ,thresholds=[0.6,0.7,0.8],
                      nms_thresholds=[0.7,0.8,0.9]):
        width,height = image.size
        min_length = min(width,height)

        min_detection_size = 12

        factor = 0.707

        scales = []
        m = min_detection_size/min_face_size
        min_length*=m

        factor_count = 0
        while min_length>min_face_size:
            scales.append((m*factor**factor_count))
            min_length*=factor
            factor_count+=1
        bounding_boxes = []
        with torch.no_grad():
            for s in scales:
                boxes = run_first_stage(image,self.pnet,scales=s,threshold=thresholds[0])
                bounding_boxes.append(boxes)

            bounding_boxes = [i for i in bounding_boxes if i is not None]
            bounding_boxes = np.vstack(bounding_boxes)

            keep = nms(bounding_boxes[:,0:5],nms_thresholds[0])
            bounding_boxes = bounding_boxes[keep]

            bounding_boxes = calibrate_box(bounding_boxes, bounding_boxes[:, 5:])
            bounding_boxes = convert_to_square(bounding_boxes)
            bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])

            img_boxes = get_image_boxes(bounding_boxes,image,size=24)
            img_boxes = torch.FloatTensor(img_boxes).to(device)

            output = self.rnet(img_boxes)
            offsets = output[0].cpu().data.numpy()
            probs = output[1].cpu().data.numpy()

            keep = np.where(probs[:,1]>thresholds[1])[0]
            bounding_boxes = bounding_boxes[keep]
            bounding_boxes[:,4] = probs[keep,1].reshape((-1,))
            offsets = offsets[keep]

            keep = nms(bounding_boxes,nms_thresholds[1])

            bounding_boxes = bounding_boxes[keep]
            bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
            bounding_boxes = convert_to_square(bounding_boxes)
            bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])

            img_boxes = get_image_boxes(bounding_boxes, image, size=48)
            if len(img_boxes) == 0:
                return [], []
            img_boxes = torch.FloatTensor(img_boxes).to(device)
            output = self.onet(img_boxes)

            landmarks = output[0].cpu().data.numpy()
            offsets = output[1].cpu().data.numpy()
            probs = output[2].cpu().data.numpy()

            keep = np.where(probs[:, 1] > thresholds[2])[0]
            bounding_boxes = bounding_boxes[keep]
            bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
            offsets = offsets[keep]
            landmarks = landmarks[keep]

            width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
            height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
            xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
            landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
            landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]

            bounding_boxes = calibrate_box(bounding_boxes, offsets)
            keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
            bounding_boxes = bounding_boxes[keep]
            landmarks = landmarks[keep]

        return bounding_boxes, landmarks

import torch
import math
from PIL import Image
import numpy as np
from .utils import _preprocess
from .nms import nms
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def run_first_stage(image, net, scale, threshold):

    width, height = image.size
    sw, sh = math.ceil(width*scale), math.ceil(height*scale)
    img = image.resize((sw, sh), Image.BILINEAR)
    img = np.asarray(img, 'float32')

    img = torch.FloatTensor(_preprocess(img)).to(device)
    with torch.no_grad():
        output = net(img)
        probs = output[1].cpu().data.numpy()[0, 1, :, :]
        offsets = output[0].cpu().data.numpy()

        boxes = _generate_bboxes(probs, offsets, scale, threshold)
        if len(boxes) == 0:
            return None
        keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
    return boxes[keep]


def _generate_bboxes(probs, offsets, scale, threshold):
    stride = 2
    cell_size = 12
    inds = np.where(probs > threshold)

    if inds[0].size == 0:
        return np.array([])

    tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
    offsets = np.array([tx1, ty1, tx2, ty2])
    score = probs[inds[0], inds[1]]

    bounding_boxes = np.vstack([
        np.round((stride*inds[1] + 1.0)/scale),
        np.round((stride*inds[0] + 1.0)/scale),
        np.round((stride*inds[1] + 1.0 + cell_size)/scale),
        np.round((stride*inds[0] + 1.0 + cell_size)/scale),
        score, offsets
    ])
    return bounding_boxes.T
权重地址：https://download.csdn.net/download/Lee_z_Adam/12469511