hed_pretrained_bsds.caffemodel
deploy.prototxt
s9xie/hed
https://github.com/s9xie/hed
https://arxiv.org/abs/1504.06375
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() |