图片边缘检测

# -*- coding: UTF-8 -*-

## https://github.com/s9xie/hed.git
## http://vcl.ucsd.edu/hed/hed_pretrained_bsds.caffemodel
## https://github.com/opencv/opencv/blob/master/samples/dnn/edge_detection.py
## https://github.com/opencv/opencv.git

import cv2
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='image.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')
args = parser.parse_args()


#! [CropLayer]
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.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]
cv2.dnn_registerLayer('Crop', CropLayer)
#! [Register]

# 模型加载
net = cv2.dnn.readNet(cv2.samples.findFile(args.prototxt), cv2.samples.findFile(args.caffemodel))

#  原始输入图像
image = cv2.imread(args.input)
(H, W) = image.shape[:2]
 
# Canny边缘检测图像
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
canny = cv2.Canny(blurred, 30, 150)


#  嵌套边缘检测图片 HED
blob = cv2.dnn.blobFromImage(image, scalefactor=1.0, size=(W, H),
    mean=(104.00698793, 116.66876762, 122.67891434),
    swapRB=False, crop=False)

net.setInput(blob)
hed = net.forward()
hed = cv2.resize(hed[0, 0], (W, H))
hed = (255 * hed).astype("uint8")
 
# show
cv2.imshow("Input", image)
cv2.imshow("Canny", canny)
cv2.imshow("HED", hed)
cv2.waitKey(0)

边缘检测