【圖像處理】pytorch篇

待整理

1. sobel算子

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms as T
from torch.autograd import Variable
import matplotlib.pyplot as plt
from PIL import Image
import cv2

class FreeDefineConV(nn.Module):
    def __init__(self, out_channel=3, in_channels=1):
        super(FreeDefineConV, self).__init__()
        kernel_x = [ [-1.0, 0.0, 1.0],
                     [-2.0, 0.0, 2.0],
                     [-1.0, 0.0, 1.0]
                     ]
        kernel_y = [[-1.0, -2.0, -1.0],
                    [ 0.0,  0.0,  0.0],
                    [-1.0,  2.0,  1.0]]

        kernel_x = torch.FloatTensor(kernel_x).expand(out_channel, in_channels, 3, 3)
        self.weight_x = nn.Parameter(data=kernel_x, requires_grad=False)

        kernel_y = torch.FloatTensor(kernel_y).t().expand(out_channel, in_channels, 3, 3)
        self.weight_y = nn.Parameter(data=kernel_y, requires_grad=False)

    def forward(self, input):
        out = F.conv2d(input, self.weight_y, stride=1,padding=0, groups=3)
        return out
Sobel_Conv = FreeDefineConV()

pil_img = Image.open('2.jpg')
tensor_img = T.ToTensor()(pil_img)

tensor_img = Sobel_Conv(tensor_img.unsqueeze(0))
print(tensor_img.max());print(tensor_img.min()) # The max and min value in tensor_img are not in range [0., 1.]
# pil_img_new = T.ToPILImage()(tensor_img.squeeze(0))
pil_img_new = tensor_img.squeeze(0).numpy()

### ------------------------------------ #
#  Normalize output into range [0., 1.]  #
# ------------------------------------ ###
pil_img_new = ( pil_img_new-pil_img_new.min() ) / ( pil_img_new.max()-pil_img_new.min() )
pil_img_new = pil_img_new.transpose(1,2,0)
# pil_img_new = cv2.cvtColor(pil_img_new, cv2.COLOR_BGR2RGB) # cv2默認爲bgr順序
pil_img_new = pil_img_new.astype('float64')


# plt.imshow(pil_img_new)
# plt.show()
cv2.namedWindow("enhanced",0)
cv2.resizeWindow("enhanced", 256, 512)
cv2.imshow("enhanced",pil_img_new)
cv2.waitKey(0)