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

#訓練集數(shù)據(jù)增強，以及將圖像數(shù)據(jù)類型轉化為張量做歸一化
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：標準差（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]
    )
])

#空間變換增強
#翻轉 → 位置不變性
#旋轉 → 方向不變性
#組合使用 → 強大的幾何不變性

#外觀變換增強
#高斯模糊 → 對圖像質量的魯棒性
#標準化

# 加載數(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)

# 隨機劃分訓練集和驗證集
indices = torch.randperm(len(dataset)).tolist()  # 隨機打亂索引
dataset_train = Subset(dataset, indices[:-valid_size])      # 訓練集：前80%
dataset_valid = Subset(dataset_test, indices[-valid_size:]) # 驗證集：后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  # 訓練時打亂
)

valid_loader = DataLoader(
    dataset_valid, batch_size=batch_size, shuffle=False  # 驗證時不打亂
)

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

# 通過命令行控制訓練輪數(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())

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

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

#總參數(shù)量: 所有參數(shù)的數(shù)量
#可訓練參數(shù): 僅計算需要梯度的參數(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ù)場景
optimizer = optim.Adam(model.parameters(), lr=lr)
# 損失函數(shù): 交叉熵損失，用于多分類
criterion = nn.CrossEntropyLoss()

# 訓練過程
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)
        # 計算損失
        loss = criterion(outputs, labels)
        train_running_loss += loss.item()
        # 計算準確率
        _, preds = torch.max(outputs.data, 1)
        train_running_correct += (preds == labels).sum().item()
        # 反向傳播
        loss.backward()
        # 更新參數(shù)
        optimizer.step()

    # 統(tǒng)計一批次的損失率和準確率
    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):
    # 設置為評估模式（禁用dropout和batchnorm的隨機性）
    model.eval()  
    print('驗證中...')
    # 累計損失值
    valid_running_loss = 0.0    
    # 累計正確預測數(shù)量
    valid_running_correct = 0   
    # 批次計數(shù)器
    counter = 0                 

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

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

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

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

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

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

    # 返回平均損失和整體準確率
    return epoch_loss, epoch_acc

# 初始化列表來跟蹤損失和準確率
# train_loss: 訓練損失列表
# valid_loss: 驗證損失列表  
# train_acc: 訓練準確率列表
# valid_acc: 驗證準確率列表
train_loss, valid_loss = [], []
train_acc, valid_acc = [], []

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

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

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

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

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

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

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

# 定義所有類別標簽的列表
# 對應模型的4個輸出類別：多云、沙漠、綠地、水域
labels = ['cloudy', 'desert', 'green_area', 'water']

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

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

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

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

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

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

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

# 對圖像應用預處理變換（調(diào)整大小、標準化等）
image = transform(image)

# 添加批次維度（將3D張量[H,W,C]轉換為4D張量[1,H,W,C]）
# 因為模型期望輸入具有批次維度
image = torch.unsqueeze(image, 0)

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

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

# 將預測的索引轉換為對應的類別名稱
pred_class = labels[int(output_label.indices)]

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

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

# 在控制臺打印真實類別和預測類別
print(f"真實類別: {gt_class}, 預測類別: {pred_class}")

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

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

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

python train.py --epochs 100

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