Skip to content

Instantly share code, notes, and snippets.

@UnaNancyOwen
Last active October 5, 2021 23:43
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save UnaNancyOwen/1dc339cbf61702ee794ecb7b2d8a3372 to your computer and use it in GitHub Desktop.
OpenCV DNN Module Edge Detection (HED) Sample
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment