import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
# 驗(yàn)證集比例
valid_split = 0.2
# 批量大小
batch_size = 64
# 數(shù)據(jù)根目錄路徑
root_dir = 'D:/pycharm projects/deep_learning_experiments/experiment_4/data'

#訓(xùn)練集數(shù)據(jù)增強(qiáng)，以及將圖像數(shù)據(jù)類型轉(zhuǎn)化為張量做歸一化
train_transform = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5)), #kernel_size:高斯卷積核大小, sigma：標(biāo)準(zhǔn)差（min,max）
    transforms.RandomRotation(degrees=(30, 70)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])
valid_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

#空間變換增強(qiáng)
#翻轉(zhuǎn) → 位置不變性
#旋轉(zhuǎn) → 方向不變性
#組合使用 → 強(qiáng)大的幾何不變性

#外觀變換增強(qiáng)
#高斯模糊 → 對(duì)圖像質(zhì)量的魯棒性
#標(biāo)準(zhǔn)化

# 加載數(shù)據(jù)集
dataset = datasets.ImageFolder(root_dir, transform=train_transform)
dataset_test = datasets.ImageFolder(root_dir, transform=valid_transform)

# 數(shù)據(jù)集信息
print(f"Classes: {dataset.classes}")  # 輸出類別名稱
dataset_size = len(dataset)
valid_size = int(valid_split * dataset_size)

# 隨機(jī)劃分訓(xùn)練集和驗(yàn)證集
indices = torch.randperm(len(dataset)).tolist()  # 隨機(jī)打亂索引
dataset_train = Subset(dataset, indices[:-valid_size])      # 訓(xùn)練集：前80%
dataset_valid = Subset(dataset_test, indices[-valid_size:]) # 驗(yàn)證集：后20%
print(f"Total training images: {len(dataset_train)}")
print(f"Total valid_images: {len(dataset_valid)}")

train_loader = DataLoader(
    dataset_train, batch_size=batch_size, shuffle=True  # 訓(xùn)練時(shí)打亂
)

valid_loader = DataLoader(
    dataset_valid, batch_size=batch_size, shuffle=False  # 驗(yàn)證時(shí)不打亂
)

import torch
import argparse
import torch.nn as nn
import torch.optim as optim
from model import build_model
from utils import save_model, save_plots
from datasets import train_loader, valid_loader, dataset
from tqdm.auto import tqdm

# 通過(guò)命令行控制訓(xùn)練輪數(shù)
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--epochs', type=int, default=20,
    help='number of epochs to train our network for')
args = vars(parser.parse_args())

# 學(xué)習(xí)參數(shù)設(shè)置
#學(xué)習(xí)率: 0.001 (Adam優(yōu)化器的常用值)
#設(shè)備: 自動(dòng)檢測(cè)GPU/CPU
lr = 0.001
epochs = args['epochs']
device = ('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Computation device: {device}\n")

# 構(gòu)建模型
model = build_model(
    pretrained=True, fine_tune=False, num_classes=len(dataset.classes)
).to(device)

#總參數(shù)量: 所有參數(shù)的數(shù)量
#可訓(xùn)練參數(shù): 僅計(jì)算需要梯度的參數(shù)
total_params = sum(p.numel() for p in model.parameters())
print(f"{total_params:,} total parameters.")
total_trainable_params = sum(
    p.numel() for p in model.parameters() if p.requires_grad)
print(f"{total_trainable_params:,} training parameters.\n")

# 優(yōu)化器: Adam，適合大多數(shù)場(chǎng)景
optimizer = optim.Adam(model.parameters(), lr=lr)
# 損失函數(shù): 交叉熵?fù)p失，用于多分類
criterion = nn.CrossEntropyLoss()

# 訓(xùn)練過(guò)程
def train(model, trainloader, optimizer, criterion):
    model.train()
    print('Training')
    train_running_loss = 0.0
    train_running_correct = 0
    counter = 0
    for i, data in tqdm(enumerate(trainloader), total=len(trainloader)):
        counter += 1
        image, labels = data
        image = image.to(device)
        labels = labels.to(device)
        # 梯度清零
        optimizer.zero_grad()
        # 前向傳播
        outputs = model(image)
        # 計(jì)算損失
        loss = criterion(outputs, labels)
        train_running_loss += loss.item()
        # 計(jì)算準(zhǔn)確率
        _, preds = torch.max(outputs.data, 1)
        train_running_correct += (preds == labels).sum().item()
        # 反向傳播
        loss.backward()
        # 更新參數(shù)
        optimizer.step()

    # 統(tǒng)計(jì)一批次的損失率和準(zhǔn)確率
    epoch_loss = train_running_loss / counter
    epoch_acc = 100. * (train_running_correct / len(trainloader.dataset))
    return epoch_loss, epoch_acc

def validate(model, testloader, criterion, class_names):
    # 設(shè)置為評(píng)估模式（禁用dropout和batchnorm的隨機(jī)性）
    model.eval()  
    print('驗(yàn)證中...')
    # 累計(jì)損失值
    valid_running_loss = 0.0    
    # 累計(jì)正確預(yù)測(cè)數(shù)量
    valid_running_correct = 0   
    # 批次計(jì)數(shù)器
    counter = 0                 

    # 初始化列表來(lái)跟蹤每個(gè)類別的準(zhǔn)確率
    # class_correct: 每個(gè)類別的正確預(yù)測(cè)數(shù)
    # class_total: 每個(gè)類別的總樣本數(shù)
    class_correct = list(0. for i in range(len(class_names)))
    class_total = list(0. for i in range(len(class_names)))

    # 禁用梯度計(jì)算以節(jié)省內(nèi)存和加速計(jì)算
    with torch.no_grad():
        # 遍歷驗(yàn)證集中的所有批次
        for i, data in tqdm(enumerate(testloader), total=len(testloader)):
            counter += 1

            # 獲取圖像數(shù)據(jù)和標(biāo)簽
            image, labels = data
            # 將數(shù)據(jù)轉(zhuǎn)移到指定設(shè)備（GPU或CPU）
            image = image.to(device)
            labels = labels.to(device)
            
            # 前向傳播：通過(guò)模型獲取預(yù)測(cè)輸出
            outputs = model(image)
            # 計(jì)算損失值
            loss = criterion(outputs, labels)
            # 累加損失值
            valid_running_loss += loss.item()
            
            # 計(jì)算準(zhǔn)確率：獲取預(yù)測(cè)類別（最大概率的索引）
            _, preds = torch.max(outputs.data, 1)
            # 累加正確預(yù)測(cè)的數(shù)量
            valid_running_correct += (preds == labels).sum().item()

            # 計(jì)算每個(gè)類別的準(zhǔn)確率
            # 獲取每個(gè)樣本的預(yù)測(cè)是否正確（布爾張量）
            correct = (preds == labels).squeeze()
            # 遍歷當(dāng)前批次中的每個(gè)樣本
            for i in range(len(preds)):
                # 獲取當(dāng)前樣本的真實(shí)標(biāo)簽
                label = labels[i]
                # 如果預(yù)測(cè)正確，則對(duì)應(yīng)類別的正確數(shù)加1
                class_correct[label] += correct[i].item()
                # 對(duì)應(yīng)類別的總樣本數(shù)加1
                class_total[label] += 1

    # 計(jì)算整個(gè)epoch的平均損失（總損失除以批次數(shù)量）
    epoch_loss = valid_running_loss / counter
    # 計(jì)算整體準(zhǔn)確率（正確預(yù)測(cè)數(shù)除以總樣本數(shù)，轉(zhuǎn)換為百分比）
    epoch_acc = 100. * (valid_running_correct / len(testloader.dataset))

    # 打印每個(gè)類別的準(zhǔn)確率
    print('\n')
    for i in range(len(class_names)):
        # 計(jì)算并打印每個(gè)類別的準(zhǔn)確率
        print(f"類別 {class_names[i]} 的準(zhǔn)確率: {100 * class_correct[i] / class_total[i]:.3f}%")
    print('\n')

    # 返回平均損失和整體準(zhǔn)確率
    return epoch_loss, epoch_acc

# 初始化列表來(lái)跟蹤損失和準(zhǔn)確率
# train_loss: 訓(xùn)練損失列表
# valid_loss: 驗(yàn)證損失列表  
# train_acc: 訓(xùn)練準(zhǔn)確率列表
# valid_acc: 驗(yàn)證準(zhǔn)確率列表
train_loss, valid_loss = [], []
train_acc, valid_acc = [], []

# 開(kāi)始訓(xùn)練循環(huán)
for epoch in range(epochs):
    # 打印當(dāng)前訓(xùn)練輪次信息
    print(f"[INFO]: 第 {epoch+1} 輪 / 共 {epochs} 輪")
    
    # 訓(xùn)練一個(gè)epoch并獲取訓(xùn)練損失和準(zhǔn)確率
    train_epoch_loss, train_epoch_acc = train(model, train_loader,
                                              optimizer, criterion)
    
    # 驗(yàn)證一個(gè)epoch并獲取驗(yàn)證損失和準(zhǔn)確率  
    valid_epoch_loss, valid_epoch_acc = validate(model, valid_loader, 
                                                 criterion, dataset.classes)
    
    # 將當(dāng)前epoch的結(jié)果添加到列表中
    train_loss.append(train_epoch_loss)
    valid_loss.append(valid_epoch_loss)
    train_acc.append(train_epoch_acc)
    valid_acc.append(valid_epoch_acc)
    
    # 打印當(dāng)前epoch的訓(xùn)練和驗(yàn)證結(jié)果
    print(f"訓(xùn)練損失: {train_epoch_loss:.3f}, 訓(xùn)練準(zhǔn)確率: {train_epoch_acc:.3f}%")
    print(f"驗(yàn)證損失: {valid_epoch_loss:.3f}, 驗(yàn)證準(zhǔn)確率: {valid_epoch_acc:.3f}%")
    print('-'*50)  # 打印分隔線
    
    # 保存訓(xùn)練好的模型權(quán)重（每個(gè)epoch都保存）
    save_model(epochs, model, optimizer, criterion)

# 訓(xùn)練完成后保存損失和準(zhǔn)確率圖表
save_plots(train_acc, valid_acc, train_loss, valid_loss)
print('訓(xùn)練完成')

# 導(dǎo)入必要的庫(kù)
import torch  # PyTorch深度學(xué)習(xí)框架
import cv2  # OpenCV計(jì)算機(jī)視覺(jué)庫(kù)，用于圖像處理
import torchvision.transforms as transforms  # PyTorch圖像變換模塊
import argparse  # 命令行參數(shù)解析庫(kù)

# 構(gòu)建參數(shù)解析器
# 創(chuàng)建ArgumentParser對(duì)象，用于處理命令行參數(shù)
parser = argparse.ArgumentParser()

# 添加輸入?yún)?shù)
# -i 或 --input: 指定輸入圖像路徑，默認(rèn)值為'input/test_data/cloudy.jpeg'
parser.add_argument('-i', '--input',
    default='input/test_data/cloudy.jpeg',
    help='輸入圖像的路徑')

# 解析輸入?yún)?shù)并將其轉(zhuǎn)換為字典格式
args = vars(parser.parse_args())

# 計(jì)算設(shè)備設(shè)置
# 指定使用CPU進(jìn)行計(jì)算（如果需要GPU可改為'cuda'）
device = 'cpu'

# 定義所有類別標(biāo)簽的列表
# 對(duì)應(yīng)模型的4個(gè)輸出類別：多云、沙漠、綠地、水域
labels = ['cloudy', 'desert', 'green_area', 'water']

# 初始化模型并加載訓(xùn)練好的權(quán)重
# pretrained=False: 不加載預(yù)訓(xùn)練權(quán)重
# fine_tune=False: 不進(jìn)行微調(diào)
# num_classes=4: 設(shè)置輸出類別數(shù)為4
model = build_model(
    pretrained=False, fine_tune=False, num_classes=4
).to(device)  # 將模型移動(dòng)到指定設(shè)備（CPU）

print('[INFO]: 正在加載自定義訓(xùn)練權(quán)重...')
# 加載訓(xùn)練好的模型檢查點(diǎn)文件
# map_location=device: 確保權(quán)重加載到正確的設(shè)備上
checkpoint = torch.load('outputs/model.pth', map_location=device)
# 將訓(xùn)練好的權(quán)重加載到模型中
model.load_state_dict(checkpoint['model_state_dict'])
# 將模型設(shè)置為評(píng)估模式（禁用dropout和batchnorm的隨機(jī)性）
model.eval()

# 定義圖像預(yù)處理變換流程
transform = transforms.Compose([
    transforms.ToPILImage(),  # 將numpy數(shù)組或tensor轉(zhuǎn)換為PIL圖像
    transforms.Resize(224),   # 調(diào)整圖像大小為224x224像素
    transforms.ToTensor(),    # 將PIL圖像轉(zhuǎn)換為tensor，并歸一化到[0,1]
    transforms.Normalize(     # 標(biāo)準(zhǔn)化處理，使用ImageNet數(shù)據(jù)集的均值和標(biāo)準(zhǔn)差
        mean=[0.485, 0.456, 0.406],  # RGB通道的均值
        std=[0.229, 0.224, 0.225]    # RGB通道的標(biāo)準(zhǔn)差
    )
])

# 讀取并預(yù)處理輸入圖像
# 使用OpenCV讀取圖像文件
image = cv2.imread(args['input'])

# 從文件路徑中提取真實(shí)類別（ground truth）
# 例如：'input/test_data/cloudy.jpeg' -> 'cloudy'
gt_class = args['input'].split('/')[-1].split('.')[0]

# 保存原始圖像的副本用于后續(xù)顯示
orig_image = image.copy()

# 將圖像從BGR格式轉(zhuǎn)換為RGB格式（OpenCV默認(rèn)使用BGR，但模型需要RGB）
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 對(duì)圖像應(yīng)用預(yù)處理變換（調(diào)整大小、標(biāo)準(zhǔn)化等）
image = transform(image)

# 添加批次維度（將3D張量[H,W,C]轉(zhuǎn)換為4D張量[1,H,W,C]）
# 因?yàn)槟Ｐ推谕斎刖哂信尉S度
image = torch.unsqueeze(image, 0)

# 在不計(jì)算梯度的情況下進(jìn)行推理（節(jié)省內(nèi)存和計(jì)算資源）
with torch.no_grad():
    # 將圖像數(shù)據(jù)移動(dòng)到設(shè)備并通過(guò)模型進(jìn)行前向傳播
    outputs = model(image.to(device))

# 獲取模型輸出的最高概率類別
# torch.topk返回前k個(gè)最大值和對(duì)應(yīng)的索引，這里k=1表示取最高概率
output_label = torch.topk(outputs, 1)

# 將預(yù)測(cè)的索引轉(zhuǎn)換為對(duì)應(yīng)的類別名稱
pred_class = labels[int(output_label.indices)]

# 在原始圖像上添加真實(shí)類別文本（綠色）
cv2.putText(orig_image,
    f"GT: {gt_class}",  # 顯示真實(shí)類別
    (10, 25),           # 文本位置坐標(biāo)（x,y）
    cv2.FONT_HERSHEY_SIMPLEX,  # 字體類型
    1,                  # 字體大小
    (0, 255, 0),        # 字體顏色（綠色）
    2,                  # 字體粗細(xì)
    cv2.LINE_AA         # 抗鋸齒線型
)

# 在原始圖像上添加預(yù)測(cè)類別文本（紅色）
cv2.putText(orig_image,
    f"Pred: {pred_class}",  # 顯示預(yù)測(cè)類別
    (10, 55),               # 文本位置坐標(biāo)（x,y），在真實(shí)類別下方
    cv2.FONT_HERSHEY_SIMPLEX,  # 字體類型
    1,                      # 字體大小
    (0, 0, 255),            # 字體顏色（紅色）
    2,                      # 字體粗細(xì)
    cv2.LINE_AA             # 抗鋸齒線型
)

# 在控制臺(tái)打印真實(shí)類別和預(yù)測(cè)類別
print(f"真實(shí)類別: {gt_class}, 預(yù)測(cè)類別: {pred_class}")

# 顯示帶有標(biāo)注的結(jié)果圖像
cv2.imshow('Result', orig_image)

# 等待按鍵操作（0表示無(wú)限等待）
cv2.waitKey(0)

# 將結(jié)果圖像保存到outputs目錄
# 文件名格式：outputs/真實(shí)類別.png
cv2.imwrite(f"outputs/{gt_class}.png", orig_image)

python train.py --epochs 100

python inference.py --input test_data/cloudy.jpg