import torch
from torch import nn
from torch.nn import functional as F

import torchvision
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
from torch.utils.data importDataLoader

import os
import numpy as np
import matplotlib.pyplot as plt

# 判斷當(dāng)前設(shè)備是否GPU
device = torch.device('cuda'if torch.cuda.is_available()else'cpu')
device

# 加載并預(yù)處理圖像
data = datasets.MNIST(root="data", 
                      train=True, 
                      transform = transforms.Compose(transforms=[transforms.ToTensor(),
                      transforms.Normalize(mean=[0.5], std=[0.5])]),
                      download=True)

# 封裝成 DataLoader
data_loader = DataLoader(dataset=data, batch_size=100, shuffle=True)

"""
    定義生成器
"""

classGenerator(nn.Module):
"""
        定義一個(gè)圖像生成
        輸入：一個(gè)向量
        輸出：一個(gè)向量（代表圖像）
    """
def__init__(self, in_features=100, out_features=28 * 28):
"""
            掛載超參數(shù)
        """
# 先初始化父類，再初始化子類
super(Generator, self).__init__()
        self.in_features = in_features
        self.out_features = out_features

# 第一個(gè)隱藏層
        self.hidden0 = nn.Linear(in_features=self.in_features, out_features=256)

# 第二個(gè)隱藏層
        self.hidden1 = nn.Linear(in_features=256, out_features=512)

# 第三個(gè)隱藏層
        self.hidden2 = nn.Linear(in_features=512, out_features=self.out_features)

defforward(self, x):

# 第一層 [b, 100] --> [b, 256]
        h = self.hidden0(x)
        h = F.leaky_relu(input=h, negative_slope=0.2)

# 第二層 [b, 256] --> [b, 512]
        h = self.hidden1(h)
        h = F.leaky_relu(input=h, negative_slope=0.2)

# 第三層 [b, 512] --> [b, 28 * 28]
        h = self.hidden2(h)

# 壓縮數(shù)據(jù)的變化范圍
        o = torch.tanh(h)

return o

"""
    定義一個(gè)鑒別器
"""

classDiscriminator(nn.Module):
"""
        本質(zhì)：二分類分類器
        輸入：一個(gè)對(duì)象
        輸出：真品還是贗品
    """
def__init__(self, in_features=28*28, out_features=1):
super(Discriminator, self).__init__()

        self.in_features=in_features
        self.out_features=out_features

# 第一個(gè)隱藏層
        self.hidden0= nn.Linear(in_features=self.in_features, out_features=512)

# 第二個(gè)隱藏層
        self.hidden1= nn.Linear(in_features=512, out_features=256)

# 第三個(gè)隱藏層
        self.hidden2= nn.Linear(in_features=256, out_features=32)

# 第四個(gè)隱藏層
        self.hidden3= nn.Linear(in_features=32, out_features=self.out_features)


defforward(self, x):

# 第一層
        h = self.hidden0(x)
        h = F.leaky_relu(input=h, negative_slope=0.2)
        h = F.dropout(input=h, p=0.2)

# 第二層
        h = self.hidden1(h)
        h = F.leaky_relu(input=h, negative_slope=0.2)
        h = F.dropout(input=h, p=0.2)

# 第三層
        h = self.hidden2(h)
        h = F.leaky_relu(input=h, negative_slope=0.2)
        h = F.dropout(input=h, p=0.2)

# 第四層
        h = self.hidden3(h)

# 輸出概率
        o = torch.sigmoid(h)

return o

"""
    構(gòu)建模型
"""
# 定義一個(gè)生成器
generator = Generator(in_features=100, out_features=784)
generator.to(device=device)

# 定義一個(gè)鑒別器
discriminator = Discriminator(in_features=784, out_features=1)
discriminator.to(device=device)

"""
    定義優(yōu)化器
"""

# 定義一個(gè)生成器的優(yōu)化器
g_optimizer = torch.optim.Adam(params=generator.parameters(), lr=1e-4)

# 定義一個(gè)鑒別的優(yōu)化器
d_optimizer = torch.optim.Adam(params=discriminator.parameters(), lr=1e-4)

"""
    定義一個(gè)損失函數(shù)
"""
loss_fn = nn.BCELoss()

# 定義訓(xùn)練輪次
num_epochs = 1000

"""
    獲取數(shù)據(jù)的標(biāo)簽
"""

defget_real_data_labels(size):
"""
        獲取真實(shí)數(shù)據(jù)的標(biāo)簽
    """
    labels = torch.ones(size,1, device=device)

return labels

defget_fake_data_labels(size):
"""
        獲取虛假數(shù)據(jù)的標(biāo)簽
    """
    labels = torch.zeros(size,1, device=device)

return labels

"""
    噪聲生成器
"""
defget_noise(size):
"""
        給生成器準(zhǔn)備數(shù)據(jù)
        - 100維度的向量
    """
    X = torch.randn(size,100, device=device)

return X

# 獲取一批測(cè)試數(shù)據(jù)

num_test_samples =16
test_noise = get_noise(num_test_samples)

"""
    訓(xùn)練過程
"""

g_losses =[]
d_losses =[]
for epoch inrange(1, num_epochs+1):

print(f"當(dāng)前正在進(jìn)行 第 {epoch} 輪 ....")

# 設(shè)置訓(xùn)練模式
    generator.train()
    discriminator.train()

# 遍歷真實(shí)的圖像
for batch_idx,(batch_real_data, _)inenumerate(data_loader):
"""
        1, 先訓(xùn)練鑒別器
            鑒別器就是一個(gè)二分類問題
            - 給一批真數(shù)據(jù)，輸出真
            - 給一批假數(shù)據(jù)，輸出假
        
        """

# 1.1 準(zhǔn)備數(shù)據(jù)
# 圖像轉(zhuǎn)向量 [b, 1, 28, 28] ---> [b, 784]
# 從數(shù)據(jù)集中獲取100個(gè)真實(shí)的手寫數(shù)字圖像
        real_data = batch_real_data.view(batch_real_data.size(0),-1).to(device=device)

# 噪聲[b, 100]
# 隨機(jī)生成100個(gè)100維度的噪聲，用于生成假圖像
        noise = get_noise(real_data.size(0))

# 根據(jù)噪聲，生成假數(shù)據(jù) 
# [b, 100] --> [b, 784]
        fake_data = generator(noise).detach()


# 1.2 訓(xùn)練過程

# 鑒別器的優(yōu)化器梯度情況
        d_optimizer.zero_grad()

# 對(duì)真實(shí)數(shù)據(jù)鑒別
        real_pred = discriminator(real_data)

# 計(jì)算真實(shí)數(shù)據(jù)的誤差
        real_loss = loss_fn(real_pred, get_real_data_labels(real_data.size(0)))

# 真實(shí)數(shù)據(jù)的梯度回傳
        real_loss.backward()


# 對(duì)假數(shù)據(jù)鑒別
        fake_pred = discriminator(fake_data)

# 計(jì)算假數(shù)據(jù)的誤差
        fake_loss = loss_fn(fake_pred, get_fake_data_labels(fake_data.size(0)))

# 假數(shù)據(jù)梯度回傳
        fake_loss.backward()

# 梯度更新
        d_optimizer.step()

# ----------------
        d_losses.append((real_loss + fake_loss).item())
# print(f"鑒別器的損失:{real_loss + fake_loss}")


"""2, 再訓(xùn)練生成器"""

# 獲取生成器的生成結(jié)果
        fake_pred = generator(get_noise(real_data.size(0)))

# 生產(chǎn)器梯度清空
        g_optimizer.zero_grad()

# 把假數(shù)據(jù)讓鑒別器鑒別一下
# 把discriminator requires_grad = False
# 設(shè)置為不可學(xué)習(xí)
for param in discriminator.parameters():
            param.requires_grad =False

        d_pred = discriminator(fake_pred)

# 設(shè)置為可學(xué)習(xí)
for param in discriminator.parameters():
            param.requires_grad =True

# 計(jì)算損失
# 把一個(gè)假東西，給專家看，專家說是真的，這個(gè)時(shí)候，造假的水平就可以了
        g_loss = loss_fn(d_pred, get_real_data_labels(d_pred.size(0)))

# 梯度回傳
        g_loss.backward()

# 參數(shù)更新
        g_optimizer.step()

# print(f"生成器誤差：{g_loss}")
        g_losses.append(g_loss.item())

# 每輪訓(xùn)練之后，觀察生成器的效果
    generator.eval()

with torch.no_grad():

# 正向推理
        img_pred = generator(test_noise)
        img_pred = img_pred.view(img_pred.size(0),28,28).cpu().data

# 畫圖
        display.clear_output(wait=True)

# 設(shè)置畫圖的大小
        fig = plt.figure(1, figsize=(12,8))
# 劃分為 4 x 4 的 網(wǎng)格
        gs = gridspec.GridSpec(4,4)

# 遍歷每一個(gè)
for i inrange(4):
for j inrange(4):
# 取每一個(gè)圖
                X = img_pred[i *4+ j,:,:]
# 添加一個(gè)對(duì)應(yīng)網(wǎng)格內(nèi)的子圖
                ax = fig.add_subplot(gs[i, j])
# 在子圖內(nèi)繪制圖像
                ax.matshow(X, cmap=plt.get_cmap("Greys"))
# ax.set_xlabel(f"{label}")
                ax.set_xticks(())
                ax.set_yticks(())
        plt.show()

import torch
from torchvision import datasets, transforms

from PIL import Image
import matplotlib.pyplot as plt

# 讀取本地下載的一張圖片
img = Image.open('girl.png')
img

resize = transforms.Resize((300, 200))
resize_img = resize(img)
resize_img

centercrop = transforms.CenterCrop(size=(200, 200))
center_img = centercrop(img)
center_img

color_jitter = transforms.ColorJitter(brightness=0.5, 
                                      contrast=0.5, 
                                      saturation=0.5, 
                                      hue=0.5)
color_jitter(img)

random_rotation = transforms.RandomRotation(degrees=10)
random_rotation(img)

from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),          # 將PIL Image轉(zhuǎn)換為Tensor
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), # 將數(shù)據(jù)歸一化到[-1, 1]之間
])