OpenCV DNN Module Edge Detection (HED) Sample

README.md

Edge Detection (Holistically-Nested Edge Detection)

• Input
  

• Output
  

Download

hed_pretrained_bsds.caffemodel
deploy.prototxt

s9xie/hed
https://github.com/s9xie/hed
https://arxiv.org/abs/1504.06375

hed.py

import os
import numpy as np
import cv2 as cv

# https://github.com/opencv/opencv/blob/4.5.3/samples/dnn/edge_detection.py#L15-L39
class CropLayer(object):
    def __init__(self, params, blobs):
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    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]]

def main():
    # キャプチャを開く
    directory = os.path.dirname(__file__)
    capture = cv.VideoCapture(os.path.join(directory, "cat.jpg")) # 画像ファイル
    #capture = cv.VideoCapture(0) # カメラ
    if not capture.isOpened():
        exit()
        
    # カスタムレイヤーを登録する
    cv.dnn_registerLayer('Crop', CropLayer)

    # モデルを読み込む
    weights = os.path.join(directory, "hed_pretrained_bsds.caffemodel")
    config = os.path.join(directory, "deploy.prototxt")
    net = cv.dnn.readNet(weights, config)

    # モデルの推論に使用するエンジンとデバイスを設定する
    net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
    net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)

    while True:
        # フレームをキャプチャして画像を読み込む
        result, image = capture.read()
        if result is False:
            cv.waitKey(0)
            break

        height, width, _ = image.shape

        # モデルの入力データを作成する
        scale = 1.0
        size = (width, height)
        mean = (104.00698793, 116.66876762, 122.67891434)
        swap = False
        crop = False
        input_blob = cv.dnn.blobFromImage(image, scale, size, mean, swap, crop)

        # モデルの入力データを設定する
        net.setInput(input_blob)

        # モデルの推論を実行する
        output_layers = net.getUnconnectedOutLayersNames()
        output_blobs = net.forward(output_layers)

        # 推論結果から画像に戻す
        edge_image = np.array(output_blobs)[5, 0, 0]
        edge_image = cv.resize(edge_image, (width, height))
        edge_image = (edge_image * 255.0).astype(np.uint8)

        # 画像を表示する
        edge_image = cv.bitwise_not(edge_image) # 白黒反転
        cv.imshow("edge detection", edge_image)
        key = cv.waitKey(10)
        if key == ord('q'):
            break

    cv.destroyAllWindows()

if __name__ == '__main__':
    main()