MTCNN工具代码

MTCNN工具代码：

import numpy as np


def iou(box, boxes, isMin=False):
    """
    :param box: [N 1] x1,y1,x2,y2 0123
    :param boxes: [N,4]-->[x1,y1,x2,y2]
    :param isMin:calculate the minimum inter area
    :return:[n, iou]
    """
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    xx1 = np.maximum(boxes[:, 0], box[0])
    yy1 = np.maximum(boxes[:, 1], box[1])
    xx2 = np.minimum(boxes[:, 2], box[2])
    yy2 = np.minimum(boxes[:, 3], box[3])
    w = np.maximum(0, xx2 - xx1)
    h = np.maximum(0, yy2 - yy1)
    inter = w * h
    if isMin:
        ious = np.true_divide(inter, np.minimum(box_area, area))
    else:
        ious = np.true_divide(inter, (box_area+area-inter))
    
    return ious


def nms(boxes, thresh=0.3, isMin=False):
    """
    :param boxes: [x1 y1 x2 y2 cls] * n cls>cls>cls --> iou -->
    :param thresh:iou condition
    :param isMin: p-r: False , o:True
    :return: np.arrays ([x,y,x1,y1,conf]* N)
    """
    if boxes.shape[0] == 0:
        return np.array([])
    a_boxes = boxes[(-boxes[:, 4]).argsort()]  # 切片和返回排序的索引按照最大到最小排列
    r_boxes = []
    while a_boxes.shape[0] > 1:  # 大于1行以上
        a_box = a_boxes[0]
        b_boxes = a_boxes[1:]
        r_boxes.append(a_box)  # 每次都保留cls最大的box
        index = np.where(iou(a_box, b_boxes, isMin) < thresh)
        # np.where 返回索引值
        a_boxes = b_boxes[index]
    x = 2
    print(x)
    if a_boxes.shape[0] > 1:
        r_boxes.append(a_boxes[0])
    
    return np.stack(r_boxes)


def nms2(boxes, thresh=0.3, isMin=False):
    """
    包含iou计算的nms，对图像中 H < 50 及 inter面积 < 5000的box进行侦测
    :param boxes: [x1 y1 x2 y2 cls] * n cls>cls>cls --> iou -->
    :param thresh:iou condition
    :param isMin: p-r: False , o:True
    :return: np.arrays ([x,y,x1,y1,conf]* N)
    """
    if boxes.shape[0] == 0:
        return np.array([])
    a_boxes = boxes[(-boxes[:, 4]).argsort()]  # 切片和返回排序的索引按照最大到最小排列
    r_boxes = []
    while a_boxes.shape[0] > 1:  # 大于1行以上
        
        a_box = a_boxes[0]
        b_boxes = a_boxes[1:]
        r_boxes.append(a_box)  # 每次都保留cls最大的box
        box_area = (a_box[2] - a_box[0]) * (a_box[3] - a_box[1])
        area = (b_boxes[:, 2] - b_boxes[:, 0]) * (b_boxes[:, 3] - b_boxes[:, 1])
        xx1 = np.maximum(b_boxes[:, 0], a_box[0])
        yy1 = np.maximum(b_boxes[:, 1], a_box[1])
        xx2 = np.minimum(b_boxes[:, 2], a_box[2])
        yy2 = np.minimum(b_boxes[:, 3], a_box[3])
        w = np.maximum(0, xx2 - xx1)
        h = np.maximum(0, yy2 - yy1)
        index_h = np.where(h < 50)
        inter = w[index_h] * h[index_h]
        index1 = np.where(inter < 5000)
        area = area[index1]
        inter = inter[index1]
        if isMin:
            ious = np.true_divide(inter, np.minimum(box_area, area))
        else:
            ious = np.true_divide(inter, (box_area + area - inter))
        index = np.where(ious < thresh)
        a_boxes = ((b_boxes[index1])[index])
    
    if a_boxes.shape[0] > 1:
        r_boxes.append(a_boxes[0])
    
    return np.stack(r_boxes)


def convert_to_squre(bbox):
    """
    :param bbox: x1,y1,x2,y2
    :return: square box of x1,y1,x2,y2
    """
    sq_box = bbox.copy()
    if bbox.shape[0] == 0:
        return np.array([])
    
    h = bbox[:, 3] - bbox[:, 1]
    w = bbox[:, 2] - bbox[:, 0]
    
    max_len = np.maximum(h, w)
    
    sq_box[:, 0] = bbox[:, 0] + (w - max_len) * 0.5
    sq_box[:, 1] = bbox[:, 1] + (h - max_len) * 0.5
    sq_box[:, 2] = sq_box[:, 0] + max_len
    sq_box[:, 3] = sq_box[:, 1] + max_len
    
    return sq_box