# -*- coding: utf-8 -*-
"""
Created on Mon Sep  3 20:24:28 2018

@author: www
"""

import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt

im = Image.open(r'C:\Users\www\Desktop\cat.png').convert('L')#讀入一張灰度圖的圖片
im = np.array(im, dtype='float32')   #將其轉化為一個矩陣

#視覺化圖片
plt.imshow(im.astype('uint8'), cmap='gray')

#將圖片矩陣轉化為pytorch tensor， 並適配卷積輸入的要求
im = torch.from_numpy(im.reshape((1, 1, im.shape[0], im.shape[1])))

#下面我們定義一個運算元對其進行輪廓檢測
#使用nn.Conv2d
conv1 = nn.Conv2d(1, 1, 3, bias=False)

#定義輪廓檢測運算元
sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32')
sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3)) #適配卷積的輸入與輸出
conv1.weight.data = torch.from_numpy(sobel_kernel) #給卷積的kernel賦值

edge1 = conv1(Variable(im)) #作用在圖片上
edge1 = edge1.data.squeeze().numpy()#將輸出轉化為圖片的格式
#視覺化邊緣檢測後的結果
plt.imshow(edge1, cmap='gray')

# 使用 F.conv2d
sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32') # 定義輪廓檢測運算元
sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3)) # 適配卷積的輸入輸出
weight = Variable(torch.from_numpy(sobel_kernel))

edge2 = F.conv2d(Variable(im), weight) # 作用在圖片上
edge2 = edge2.data.squeeze().numpy() # 將輸出轉換為圖片的格式
plt.imshow(edge2, cmap='gray')

#可以看到兩種形式能夠得到相同的效果，不同的地方相信你也看到了，使用 nn.Conv2d() 相當於直接定義了一層卷積網路結構，而使用 
#torch.nn.functional.conv2d() 相當於定義了一個卷積的操作，所以使用後者需要再額外去定義一個 weight，而且這個 weight 
#也必須是一個 Variable，而使用 nn.Conv2d() 則會幫我們預設定義一個隨機初始化的 weight，如果我們需要修改，那麼取出其中
#的值對其修改，如果不想修改，那麼可以直接使用這個預設初始化的值，非常方便

#實際使用中我們基本都使用 nn.Conv2d() 這種形式


#池化層
#卷積網路中另外一個非常重要的結構就是池化，這是利用了圖片的下采樣不變性，即一張圖片變小了還是能夠看出了這張圖片的內容，
#而使用池化層能夠將圖片大小降低，非常好地提高了計算效率，同時池化層也沒有引數。池化的方式有很多種，比如最大值池化，均值
#池化等等，在卷積網路中一般使用最大值池化。

#在 pytorch 中最大值池化的方式也有兩種，一種是 nn.MaxPool2d()，一種是 torch.nn.functional.max_pool2d()，
#他們對於圖片的輸入要求跟卷積對於圖片的輸入要求是一樣了，就不再贅述

# 使用 nn.MaxPool2d
pool1 = nn.MaxPool2d(2, 2)
print('before max pool, image shape: {} x {}'.format(im.shape[2], im.shape[3]))
small_im1 = pool1(Variable(im))
small_im1 = small_im1.data.squeeze().numpy()
print('after max pool, image shape: {} x {} '.format(small_im1.shape[0], small_im1.shape[1]))

# F.max_pool2d
print('before max pool, image shape: {} x {}'.format(im.shape[2], im.shape[3]))
small_im2 = F.max_pool2d(Variable(im), 2, 2)
small_im2 = small_im2.data.squeeze().numpy()
print('after max pool, image shape: {} x {} '.format(small_im1.shape[0], small_im1.shape[1]))
plt.imshow(small_im2, cmap='gray')