使用libtorch实现sobel算子进行边缘检测，python和c++实现

python

import torch
import torch.nn as nn
from PIL import Image
import numpy as np
import cv2


class SobelEdgeDetector(nn.Module):
    def __init__(self):
        super(SobelEdgeDetector, self).__init__()
        self.convx = nn.Conv2d(1,1,3,padding=1)
        self.convy = nn.Conv2d(1,1,3,padding=1)

        self.convx.weight.data = torch.FloatTensor([[1,0,-1],[2,0,-2],[1,0,-1]]).unsqueeze(0).unsqueeze(0)
        self.convy.weight.data = torch.FloatTensor([[1,2,1],[0,0,0],[-1,-2,-1]]).unsqueeze(0).unsqueeze(0)
    
    def forward(self, x):
        grad_x = self.convx(x)
        grad_y = self.convy(x)

        edges = torch.sqrt(grad_x**2 + grad_y**2)

        return edges
    
def norm_gray(img):
    max_value = np.max(img)
    min_value = np.min(img)

    img = ( img - min_value) / (max_value - min_value ) * 255
    img = np.clip(img, 0, 255).astype(np.ubyte)

    return img
    


if __name__ == '__main__':
    img = cv2.imread("/home/lhb_python_2/seg_2d/JPEGImages/Abyssinian_1.jpg", 0)
    img = torch.tensor(img,dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    print(img.shape)

    model = SobelEdgeDetector()
    model.eval()
    edges = model(img)
    outputs = edges.squeeze().detach().numpy()
    print(outputs.shape)
    
    outputs = norm_gray(outputs)

    cv2.imwrite("edge.jpg",outputs)

python结果

cpp

// libtorch实现sobel算子边缘检测

#include
#include
#include
#include

using std::cout; using std::endl;
using std::string;


inline torch::nn::Conv2dOptions Conv2dOptions(
	int64_t in_channels, int64_t out_channles, int64_t kernel_size,
	int64_t stride = 1, int64_t padding = 1, bool with_bias = false
) 
{
	torch::nn::Conv2dOptions conv2dOptions = torch::nn::Conv2dOptions(in_channels, out_channles, kernel_size);
	conv2dOptions.stride(stride);
	conv2dOptions.padding(padding);
	conv2dOptions.bias(with_bias);
	
	return conv2dOptions;
}


class SobelEdgeDetector : public torch::nn::Module
{
private:
	torch::nn::Conv2d convx{ nullptr };
	torch::nn::Conv2d convy{ nullptr };

public:
	SobelEdgeDetector(int in_channels, int out_channels, int kernel_size, int stride, int padding = 1);
	torch::Tensor forward(torch::Tensor x);
};



int main()
{
	SobelEdgeDetector model(1, 1, 3, 1, 1);

	std::string imgPath("C:\\Users\\Administrator\\Downloads\\Abyssinian_1.jpg");
	cv::Mat img = cv::imread(imgPath, 0);
	torch::Tensor imgTesor = torch::from_blob(img.data, { img.rows, img.cols, 1 }, torch::kByte).permute({ 2,0,1 }).unsqueeze(0).to(torch::kFloat32);
	cout << imgTesor.sizes() << endl;
	
	model.eval();
	auto outTensor = model.forward(imgTesor);
	//cout << outTensor.sizes() << endl;
	//cout << outTensor << endl;

	
	float min_val = outTensor.min().item<float>();
	float max_val = outTensor.max().item<float>();
	outTensor = (outTensor - min_val) / (max_val - min_val) * 255;
	

	cv::Mat outImg(outTensor.size(2), outTensor.size(3), CV_32FC1, outTensor.data_ptr<float>());

	/*cv::imshow("w", outImg);
	cv::waitKey();*/
	cv::imwrite("edge.jpg", outImg);

	return 0;
}

SobelEdgeDetector::SobelEdgeDetector(int in_channels, int out_channels, int kernel_size, int stride, int padding)
{
	convx = register_module("convx", torch::nn::Conv2d(Conv2dOptions(in_channels, out_channels, kernel_size, stride, padding)));
	convy = register_module("convy", torch::nn::Conv2d(Conv2dOptions(in_channels, out_channels, kernel_size, stride, padding)));

	convx->weight = torch::tensor({ {1,0,-1},{2,0,-2},{1,0,-1 } }, torch::kFloat32);
	convx->weight = convx->weight.unsqueeze(0);
	convx->weight = convx->weight.unsqueeze(0);
	//cout << convx->weight << endl;

	convy->weight = torch::tensor({ {1,2,1},{0,0,0},{-1,-2,-1 } }, torch::kFloat32);
	convy->weight = convy->weight.unsqueeze(0);
	convy->weight = convy->weight.unsqueeze(0);
}

torch::Tensor SobelEdgeDetector::forward(torch::Tensor x)
{
	//return torch::Tensor();

	torch::Tensor gradX = this->convx(x);
	torch::Tensor gradY = this->convy(x);
	torch::Tensor edges = torch::sqrt(gradX * gradX + gradY * gradY);
	
	return edges;
}

c++结果