使用pytorch卷积图像

使用pytorch卷积图像：通常神经网络中的参数是可以学习的，当然我们也可以设置固定的网络参数（如卷积核kernel）来卷积图像，下面是一个示例，示例的步骤：

1. 使用opencv读取图像，并且交换维度和pytorch的数据存储维度NCHW保持一致（opencv是HWC）

2. 使用batch size为2进行卷积，所以输出是两个通道的图像

3. inputs维度为：2x3x512x512(NCHW)； 卷积维度：1x3x3x3(CoCiKhKw)；outputs维度为：2x1x512x512(NCWH)

import torch.nn as nn
import torch
import cv2
import numpy as np


if __name__ == "__main__":
    img = cv2.imread('/home/wang/lena.jpg')
    img = np.swapaxes(img, 0, 2)
    img = np.swapaxes(img, 1, 2).reshape(3, 512, 512)
    img = img.astype(np.float32) / 255.0
    inputs = torch.from_numpy(np.array([img, img]))
    inputs = inputs.cuda()
    print(inputs.shape)

    conv = nn.Conv2d(3, 1, kernel_size=3, padding=1, groups=1, stride=1, bias=False)
    a = [[1, 1, 1], [1, -8, 1], [1, 1, 1]]
    conv.weight.data = torch.Tensor([a, a, a]).unsqueeze(0).cuda()
    print(conv.weight.shape)

    outputs = conv(inputs)

    outputs = outputs.cpu()
    print(outputs.shape)
    img = outputs.detach().numpy()

    img = np.clip(img, 0, 100)
    img = 255 * ((img - img.min()) / (img.max() - img.min()))
    img1 = img[0, 0, :, :]
    img2 = img[1, 0, :, :]
    cv2.imwrite("/home/wang/xxx.jpg", img1.astype(np.uint8))
    cv2.imwrite("/home/wang/yyy.jpg", img2.astype(np.uint8))
    # cv2.imshow("window", img.astype(np.uint8))

 下图分别是lena.jpg, xxx.jpg 和 yyy.jpg