类封装Canny算法

1. 代码

import cv2 as cv
import numpy as np
import math
import time
import torch.nn.functional as F
import torch


def show_matrix(_message, _matrix, shape_sign=False):
	"""

	:param _message: window-frame name
	:param _matrix: matrix of image
	:param shape_sign: if True, show the shape of matrix
	:return: none
	"""
	if shape_sign:
		print(_message, "shape: ", _matrix.shape)
	cv.imshow(_message, _matrix)
	exit_button = cv.waitKey(0)  # increase waitKey will make motion delay
	if exit_button == 27:
		cv.destroyAllWindows()


def show_cost(a, b, message):
	print(message, ':{0}'.format(b - a))


def conv(_input, _filter):
	result_h = _input.shape[0] - _filter.shape[0] + 1
	result_w = _input.shape[1] - _filter.shape[1] + 1
	_result = np.zeros([result_h, result_w])
	for r in range(0, result_h):
		for c in range(0, result_w):
			cur_input = _input[r:r + _filter.shape[0], c:c + _filter.shape[1]]
			cur_output = cur_input * _filter
			conv_sum = np.sum(cur_output)
			_result[r, c] = conv_sum
	return _result


def canny_func(_img, threshold1, threshold2):
	# gaussian
	st = time.time()
	gray = cv.cvtColor(_img, cv.COLOR_BGR2GRAY)
	gaussian_blur = cv.GaussianBlur(gray, (5, 5), 1)
	gaussian_blur_show = np.uint8(np.copy(gaussian_blur))
	show_matrix('gaussian', gaussian_blur_show)
	show_cost(st, time.time(), 'gaussian')

	# gradient
	st = time.time()
	# width, height = gaussian_blur.shape[:2]
	# dx = np.zeros([width - 2, height - 2])
	# dy = np.zeros([width - 2, height - 2])
	# d = np.zeros([width - 2, height - 2])
	# d_degree = np.zeros([width - 2, height - 2])

	sobel_x = np.array([1, 2, 1, 0, 0, 0, -1, -2, -1]).reshape(3, 3)
	sobel_y = np.array([1, 0, -1, 2, 0, -2, 1, 0, -1]).reshape(3, 3)

	at = torch.from_numpy(gaussian_blur.astype(np.long)).unsqueeze(0).unsqueeze(0)
	bxt = torch.from_numpy(sobel_x.astype(np.long)).unsqueeze(0).unsqueeze(0)
	byt = torch.from_numpy(sobel_y.astype(np.long)).unsqueeze(0).unsqueeze(0)
	dx = F.conv2d(at, bxt, padding=0, groups=1).squeeze(0).squeeze(0).numpy().astype(np.float64)
	dy = F.conv2d(at, byt, padding=0, groups=1).squeeze(0).squeeze(0).numpy().astype(np.float64)

	# dx = cv.filter2D(gaussian_blur, 2, sobel_x.T).astype(np.float64)		# cv卷积结果不对，不知道为什么
	# dy = cv.filter2D(gaussian_blur, 2, sobel_y.T).astype(np.float64)

	# dx = conv(gaussian_blur, sobel_x)					# 方法2,耗时
	# dy = conv(gaussian_blur, sobel_y)
	d = np.sqrt(np.square(dx) + np.square(dy))
	d_degree = np.arctan2(dy, dx) * 180 / math.pi		# rad2deg
	d_degree[d_degre