python

import torch
import torch.nn as nn
from PIL import Image
import numpy as np
import cv2class 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 edgesdef norm_gray(img):max_value = np.max(img)min_value = np.min(img)img = ( img - min_value) / (max_value - min_value ) * 255img = np.clip(img, 0, 255).astype(np.ubyte)return imgif __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<torch/script.h>
#include<torch/torch.h>
#include<opencv.hpp>
#include<iostream>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++结果