python+opencv调用HED(Holistically-Nested Edge Detection)模型进行边缘检测

       opencv3.4以来逐渐加大了对Dnn模块的支持,在最新的opencv中支持了对HED模型的调用。

HED模型出自论文,Holistically-Nested Edge Detection ,ICCV2015,Marr奖提名,非常值得看。 https://arxiv.org/abs/1504.06375

      采用opencv的Dnn模块调用训练好的HED模型需要先下载模型文件http://vcl.ucsd.edu/hed/hed_pretrained_bsds.caffemodel

和deploy.prototxthttps://github.com/s9xie/hed/tree/master/examples/hed

下面附上python调用代码(来自opencv例程,修改了错误部分)

  

import cv2 as cv
import argparse

parser = argparse.ArgumentParser(
        description='This sample shows how to define custom OpenCV deep learning layers in Python. '
                    'Holistically-Nested Edge Detection (https://arxiv.org/abs/1504.06375) neural network '
                    'is used as an example model. Find a pre-trained model at https://github.com/s9xie/hed.')
parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera', default='1.jpg')
parser.add_argument('--prototxt', help='Path to deploy.prototxt', default='deploy.prototxt')
parser.add_argument('--caffemodel', help='Path to hed_pretrained_bsds.caffemodel', default='hed_pretrained_bsds.caffemodel')
parser.add_argument('--width', help='Resize input image to a specific width', default=500, type=int)
parser.add_argument('--height', help='Resize input image to a specific height', default=500, type=int)
args = parser.parse_args()

#! [CropLayenr]
class CropLayer(object):
    def __init__(self, params, blobs):
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    # Our layer receives two inputs. We need to crop the first input blob
    # to match a shape of the second one (keeping batch size and number of channels)
    def getMemoryShapes(self, inputs):
        inputShape, targetShape = inputs[0], inputs[1]
        batchSize, numChannels = inputShape[0], inputShape[1]
        height, width = targetShape[2], targetShape[3]

        #self.ystart = (inputShape[2] - targetShape[2]) / 2
        #self.xstart = (inputShape[3] - targetShape[3]) / 2


        self.ystart = int((inputShape[2] - targetShape[2]) / 2)
        self.xstart = int((inputShape[3] - targetShape[3]) / 2)

        self.yend = self.ystart + height
        self.xend = self.xstart + width

        return [[batchSize, numChannels, height, width]]

    def forward(self, inputs):
        return [inputs[0][:,:,self.ystart:self.yend,self.xstart:self.xend]]
#! [CropLayer]

#! [Register]
cv.dnn_registerLayer('Crop', CropLayer)
#! [Register]

# Load the model.
net = cv.dnn.readNet(cv.samples.findFile(args.prototxt), cv.samples.findFile(args.caffemodel))

kWinName = 'Holistically-Nested Edge Detection'
cv.namedWindow('Input', cv.WINDOW_NORMAL)
cv.namedWindow(kWinName, cv.WINDOW_NORMAL)


frame=cv.imread('5.jpg')


cv.imshow('Input', frame)
#cv.waitKey(0)

inp = cv.dnn.blobFromImage(frame, scalefactor=1.0, size=(args.width, args.height),
                               mean=(104.00698793, 116.66876762, 122.67891434),
                               swapRB=False, crop=False)
net.setInput(inp)

out = net.forward()
out = out[0, 0]
out = cv.resize(out, (frame.shape[1], frame.shape[0]))
cv.imshow(kWinName, out)
cv.imwrite('result.png',out)
cv.waitKey(0)

效果图:

python+opencv调用HED(Holistically-Nested Edge Detection)模型进行边缘检测_第1张图片
 

你可能感兴趣的:(opencv图像处理)