MTCNN对O网络的测试

简易版

cpimg = img
    imgdraw = ImageDraw.Draw(img)
    imglist = []
    for box in boxs:
        cropimg = cpimg.crop((box[0], box[1], box[2], box[3]))
        cropimg = cropimg.resize((48, 48), Image.ANTIALIAS)
        cropimg = np.array(cropimg)
        cropimg = np.transpose(cropimg, [2, 0, 1])
        imglist.append(cropimg)
    imglist = np.array(imglist)
    imglist = torch.FloatTensor(imglist)
    if torch.cuda.is_available():
        imglist = imglist.cuda()
    confidence, offset = net(imglist)
    confidence = confidence.cpu().data.numpy()
    offset = offset.cpu().data.numpy()
    confidence = confidence[:, 0, 0, 0]
    offset = offset[:, :, 0, 0]
    indexs = np.where(confidence > 0.7)
    conf = confidence[indexs]
    off = offset[indexs]
    boxs = np.array(boxs)
    bboxs = boxs[indexs]
    bboxs = NmsDo(bboxs, "UNIU")
    print(len(bboxs))
    for box in bboxs:
        imgdraw.rectangle((box[0], box[1], box[2], box[3]))
    img.show()
    return bboxs

 

imglist = []
    oboxslist = []
    boxss = np.array(boxs)
    if boxss.shape[0] == 0:
        return []

    #把框变成48*48的正方形
    x1 = boxss[:, 0]
    y1 = boxss[:, 1]
    x2 = boxss[:, 2]
    y2 = boxss[:, 3]
    w = x2 - x1
    h = y2 - y1
    sidelen = np.maximum(w, h)
    cx = x1 + w / 2
    cy = y1 + h / 2
    _x1 = cx - sidelen / 2
    _y1 = cy - sidelen / 2
    _x2 = cx + sidelen / 2
    _y2 = cy + sidelen / 2

    _boxs = np.stack([_x1, _y1, _x2, _y2], axis=1)
    for box in _boxs:
        cropimg = img.crop((box[0], box[1], box[2], box[3]))
        imgdata = cropimg.resize((48, 48), Image.ANTIALIAS)
        cropimg = np.array(imgdata, dtype=np.float32)/255.
        cropimg = cropimg.transpose([2, 0, 1]) #CHW
        imglist.append(cropimg)

    #传入O网络
    imgdata = np.array(imglist)
    imgdata = torch.FloatTensor(imgdata)
    if torch.cuda.is_available():
        imgdata = imgdata.cuda()

    if net.island:
        confidence, offset, landoff = net(imgdata)
        landoff = landoff.view(-1, 10)  # N * 10
        landoff = landoff.cpu().data.numpy()
    else:
        confidence, offset = net(imgdata)

    confidence = confidence.view(-1, 1) #N * 1
    offset = offset.view(-1, 4) #N * 4
    confidence = confidence.cpu().data.numpy()
    offset = offset.cpu().data.numpy()
    # 筛选置信度大于阈值的框
    if show_conf:
        BoxsAndConfidence(img, _boxs, confidence, offset)
    indexs = np.where(confidence > 0.9)  # np.where()[0] 表示行的索引 1是列
    indexs = np.stack(indexs, axis=1)
    if indexs.shape[0] > 0:
        # 反算坐标
        # for index in indexs:
        # 建议框坐标
        _x1 = _boxs[indexs[:,0],0]
        _y1 = _boxs[indexs[:,0],1]
        _x2 = _boxs[indexs[:,0],2]
        _y2 = _boxs[indexs[:,0],3]
        w = _x2 - _x1
        h = _y2 - _y1
        offx1 = offset[indexs[:,0],0]
        offy1 = offset[indexs[:,0],1]
        offx2 = offset[indexs[:,0],2]
        offy2 = offset[indexs[:,0],3]
        # 真实框坐标
        x1 = _x1 + offx1 * w
        y1 = _y1 + offy1 * h
        x2 = _x2 + offx2 * w
        y2 = _y2 + offy2 * h
        conf = confidence[indexs[:,0], indexs[:,1]]
        #5个特征点的坐标
        if net.island:
            landoffx1 = landoff[indexs[:,0],0]
            landoffy1 = landoff[indexs[:,0],1]
            landoffx2 = landoff[indexs[:,0],2]
            landoffy2 = landoff[indexs[:,0],3]
            landoffx3 = landoff[indexs[:,0],4]
            landoffy3 = landoff[indexs[:,0],5]
            landoffx4 = landoff[indexs[:,0],6]
            landoffy4 = landoff[indexs[:,0],7]
            landoffx5 = landoff[indexs[:,0],8]
            landoffy5 = landoff[indexs[:,0],9]

            fx1 = _x1 + landoffx1 * w
            fy1 = _y1 + landoffy1 * h
            fx2 = _x1 + landoffx2 * w
            fy2 = _y1 + landoffy2 * h
            fx3 = _x1 + landoffx3 * w
            fy3 = _y1 + landoffy3 * h
            fx4 = _x1 + landoffx4 * w
            fy4 = _y1 + landoffy4 * h
            fx5 = _x1 + landoffx5 * w
            fy5 = _y1 + landoffy5 * h

            oboxslist.extend(np.stack([x1, y1, x2, y2, conf
                            , fx1, fy1, fx2, fy2, fx3, fy3, fx4, fy4, fx5, fy5],axis=1))
        else:
            oboxslist.extend(np.stack([x1, y1, x2, y2, conf],axis=1))

    #做NMS
    oklist = NmsDo(oboxslist, "OTHER")
    if imgshow == True:
        h_img = ImageDraw.Draw(img)
        if net.island:
            for box in oklist:
                h_img.rectangle((box[0], box[1], box[2], box[3]), outline="red")
                h_img.point([(box[5], box[6]),(box[7], box[8]),(box[9], box[10]),(box[11], box[12]),(box[13], box[14])]
                            ,fill=(255, 0, 0))
        else:
            for box in oklist:
                h_img.rectangle((box[0], box[1], box[2], box[3]), outline="red")

        if not isuse:
            img.show()
    return oklist