MTCNN对P网络的测试

def P(net,img):
    #获取原始图片的宽和高
    width,high = img.size
    #缩放比例
    scale = 1
    #复制原始图片以做后续处理
    cpimg = img
    #创建一个列表用于保存一系列的框
    boxlist = []
    #获取图片宽和高的最小值作为结束循环的条件
    sidelen = np.minimum(width,high)
    #循环构造建议框直到图片尺寸小于12
    while sidelen > 12:
        #对图片进行归一化
        imgdata = np.array(cpimg,dtype=np.float32)/255
        #对图片进行升维
        imgdata = np.expand_dims(imgdata,axis=0)
        #变成NCHW的形式
        imgdata = np.transpose(imgdata,[0,3,1,2])
        #变成Tensor
        imgdata = torch.FloatTensor(imgdata)
        if torch.cuda.is_available():
            imgdata = imgdata.cuda()
        #将图片传入网络获得置信度和偏移
        confidence,offset = net(imgdata)
        #获取真正的置信度和偏移
        confidence = confidence[0][0].cpu().data.numpy()
        offset = offset[0].cpu().data.numpy()
        #得到对应的坐标偏移
        off_x1 = offset[0]
        off_y1 = offset[1]
        off_x2 = offset[2]
        off_y2 = offset[3]
        #得到置信度大于0.6的像素值的坐标
        indexs = np.where(confidence > 0.6)
        #将对应的x和y进行组合，y在前x在后
        indexs = np.stack(indexs,axis=1)
        #得到建议框的坐标
        _x1 = (indexs[:,1]*2)/scale
        _y1 = (indexs[:,0]*2)/scale
        _x2 = (indexs[:,1]*2+12)/scale
        _y2 = (indexs[:,0]*2+12)/scale
        #得到建议框的宽和高
        w = _x2 - _x1
        h = _y2 - _y1
        #得到建议框坐标所对应的坐标偏移
        offx1 = off_x1[indexs[:,0],indexs[:,1]]
        offy1 = off_y1[indexs[:,0],indexs[:,1]]
        offx2 = off_x2[indexs[:,0],indexs[:,1]]
        offy2 = off_y2[indexs[:,0],indexs[:,1]]
        #得到可能是人脸的像素的坐标的置信度
        conf = confidence[indexs[:,0],indexs[:,1]]
        #实际框的坐标
        x1 = _x1 + offx1*w
        y1 = _y1 + offy1*h
        x2 = _x2 + offx2*w
        y2 = _y2 + offy2*h
        #将这个框及其置信度添加进框列表中
        boxlist.extend(np.stack([x1,y1,x2,y2,conf],axis=1))
        #对图片进行缩放
        #让缩放因子进行累乘以为图像金字塔
        scale *= 0.7
        #得到缩放后的宽和高
        _width = width*scale
        _high = high*scale
        #将图片缩放成缩放后的宽和高
        cpimg = cpimg.resize((_width,_high),Image.ANTIALIAS)
        #获取缩放后的图片的宽和高
        sidelen = np.minimum(_width,_high)
    #对这堆框做非极大抑制获取置信度最大的框
    bboxs = NmsDo(boxlist,mode="UNIUM")
    #画出这些框
    imgdraw = ImageDraw.Draw(img)
    for box in bboxs:
        imgdraw.rectangle((box[0],box[1],box[2],box[3]),outline='red')
        img.show()
    return bboxs

def PnetDetect(net, img, imgshow=False):
    width, high = img.size
    scale = 1  # 缩放比列
    copyimg = img
    boxslist = []  # 框列表
    sidelen = np.minimum(width, high)
    while sidelen > 12:  # 最小尺寸
        imgdata = np.array(copyimg, dtype=np.float32) / 255
        imgdata = np.expand_dims(imgdata, axis=0)  # 扩展维度变成4维
        imgdata = imgdata.transpose([0, 3, 1, 2])  # 轴变换NCHW
        imgdata = torch.FloatTensor(imgdata)
        if torch.cuda.is_available():
            imgdata = imgdata.cuda()
        confidence, offset = net(imgdata)  # 置信度和偏移
        confidence = confidence[0][0].cpu().data.numpy()
        offset = offset[0].cpu().data.numpy()
        # confidence = np.squeeze(confidence)  # 降维(h, w) ,去掉维度为1的
        # offset = np.squeeze(offset)  # 降维(c, h, w)
        off_x1 = offset[0]
        off_y1 = offset[1]
        off_x2 = offset[2]
        off_y2 = offset[3]
        # 筛选置信度大于阈值的框，筛选特征层上置信度大于0.6的像素
        indexs = np.where(confidence > 0.6)  # np.where()[0] 表示行的索引 1是列
        indexs = np.stack(indexs, axis=1)#将index在维度为1的地方合在一起，就得到了需要的置信度的坐标
        # print(confidence[indexs[:,0],indexs[:,1]])#打印需要的置信度
        if indexs.shape[0] > 0:
            # 反算坐标
            # for index in indexs:
            _x1 = (indexs[:, 1] * 2) / scale
            _y1 = (indexs[:, 0] * 2) / scale
            _x2 = (indexs[:, 1] * 2 + 12) / scale
            _y2 = (indexs[:, 0] * 2 + 12) / scale
            w = _x2 - _x1
            h = _y2 - _y1
            offx1 = off_x1[indexs[:, 0], indexs[:, 1]]
            offy1 = off_y1[indexs[:, 0], indexs[:, 1]]
            offx2 = off_x2[indexs[:, 0], indexs[:, 1]]
            offy2 = off_y2[indexs[:, 0], indexs[:, 1]]
            conf = confidence[indexs[:, 0], indexs[:, 1]]
            # 真实框坐标
            x1 = _x1 + offx1 * w
            y1 = _y1 + offy1 * h
            x2 = _x2 + offx2 * w
            y2 = _y2 + offy2 * h
            # h_img.rectangle((x1, y1, x2, y2), outline="red")
            boxslist.extend(np.stack([x1, y1, x2, y2, conf], axis=1))

        # 图片缩放
        scale *= 0.7#缩放因子进行累乘以做图像金字塔
        _width = int(width * scale)
        _high = int(high * scale)
        copyimg = copyimg.resize((_width, _high), Image.ANTIALIAS)
        sidelen = np.minimum(_width, _high)
    #NMS
    oklist = NmsDo(boxslist, "UNIUM")
    if imgshow == True:
        h_img = ImageDraw.Draw(img)
        for box in oklist:
            h_img.rectangle((box[0],box[1],box[2],box[3]), outline="red")
            h_img.text((box[0],box[1]), str(box[4]), "black")
        img.show()
    return oklist