import os
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

class CatDogDataset(Dataset):
    """
    自定义的猫狗二分类数据集，继承自 PyTorch 的 Dataset 类。
    支持训练与测试模式，自动读取图像路径并生成标签。

    Args:
        root_dir (str): 数据集根目录，比如 'data/train' 或 'data/test'
        transform (callable, optional): 数据增强与预处理操作
    """

    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir              # 存储 train 或 test 路径
        self.transform = transform            # 图像预处理 transform

        # label 映射：猫为 0，狗为 1
        self.label_map = {'cat': 0, 'dog': 1}

        self.image_paths = []                 # 保存所有图像的路径
        self.labels = []                      # 保存对应标签

        # 遍历每个类别文件夹
        for label_name in ['cat', 'dog']:
            class_dir = os.path.join(root_dir, label_name)
            for filename in os.listdir(class_dir):
                if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
                    self.image_paths.append(os.path.join(class_dir, filename))
                    self.labels.append(self.label_map[label_name])

    def __len__(self):
        """返回数据集大小"""
        return len(self.image_paths)

    def __getitem__(self, idx):
        """
        根据索引 idx 返回一个样本（图像, 标签）
        """
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        # 打开图像
        image = Image.open(img_path).convert('RGB')  # 保证统一为 RGB 格式

        # 应用预处理/增强
        if self.transform:
            image = self.transform(image)

        return image, label

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

class SimpleCNN(nn.Module):
    """
    一个用于图像二分类的简单卷积神经网络（CNN）模型。
    结构：Conv + ReLU + Pooling x3 + FC x2
    输入图像默认尺寸为 3x224x224
    """

    def __init__(self, num_classes=2):
        """
        初始化模型结构

        Args:
            num_classes (int): 输出类别数，默认是2（cat, dog）
        """
        super(SimpleCNN, self).__init__()

        # 卷积层1：输入通道3，输出通道16，核3x3
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        # 卷积层2
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        # 卷积层3
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)

        # 最大池化层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # 全连接层1：64个通道，图像缩小3次 => 224 -> 112 -> 56 -> 28 -> flatten: 64 * 28 * 28
        self.fc1 = nn.Linear(64 * 28 * 28, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        """
        定义前向传播流程

        Args:
            x (Tensor): 输入图像张量，形状为 (N, 3, 224, 224)

        Returns:
            Tensor: 输出分类 logits，形状为 (N, num_classes)
        """
        x = self.pool(F.relu(self.conv1(x)))  # -> (N, 16, 112, 112)
        x = self.pool(F.relu(self.conv2(x)))  # -> (N, 32, 56, 56)
        x = self.pool(F.relu(self.conv3(x)))  # -> (N, 64, 28, 28)

        x = x.view(x.size(0), -1)  # 展平为向量

        x = F.relu(self.fc1(x))    # -> (N, 128)
        x = self.fc2(x)            # -> (N, num_classes)

        return x

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from dataset import CatDogDataset
from model import SimpleCNN
import torchvision.transforms as transforms
from tqdm import tqdm  

# 使用 GPU（如果可用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"当前使用设备: {device}")

# 参数配置
EPOCHS = 10
BATCH_SIZE = 32
LR = 0.001
IMG_SIZE = 224
SAVE_PATH = 'best_model.pth'

# 数据预处理
train_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.RandomHorizontalFlip(),  
    # 数据增强，在训练过程中，以 50% 的概率 将图像 水平翻转（即左右对调）。
    # 增加样本多样性，让模型见到更多“不同角度”的数据；提高泛化能力，减少过拟合，学到“猫是什么”而不是“猫朝哪边”；模拟现实场景	实际拍摄中物体方向可能左右都有，训练集中未必全有覆盖
    # transforms.RandomHorizontalFlip(p=0.7)70% 概率翻转。
    # 在猫狗分类、人脸检测中可以用，但是含有文字，或者方位信息重要的（道路场景）不能用。
    # 还可以transforms.RandomRotation(10)旋转增强
    transforms.ToTensor(), # 把图像从 [0, 255] 映射到 [0.0, 1.0]
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3) 
    # 把每个像素值归一化到一个相对“标准”的分布范围，便于神经网络更快、更稳定地训练。
])

test_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])

# 加载数据
train_dataset = CatDogDataset(root_dir='data/train', transform=train_transform)
test_dataset = CatDogDataset(root_dir='data/test', transform=test_transform)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
# shuffle 防止模型记住顺序（比如训练集如果前半全是猫，后半全是狗，模型可能学到“顺序”，而不是图像特征），提高泛化能力（打乱样本分布后，能更好地逼近整体数据的真实分布），提升收敛稳定性（每轮训练样本顺序都不一样，能避免局部最小值或震荡问题）
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
# 测试集是用来“评估模型效果”的，必须稳定、可重复。如果你在测试时也打乱数据顺序，会导致：测试结果不稳定（每次测试顺序变了）；模型对样本的预测顺序也变了，难以做准确率等指标对比；可视化结果无法复现

# 构建模型
model = SimpleCNN(num_classes=2).to(device)

# 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)
# 优化器（Optimizer）决定了模型如何更新参数，让损失函数变小，从而让模型更聪明。优化器就是根据损失函数的梯度信息，去更新模型参数（如权重和偏置） 的算法。
# Adam自适应学习率的优化器，收敛快、效果好，最常用

# 训练函数
def train_one_epoch(model, loader, optimizer, criterion):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    pbar = tqdm(loader, desc='Training', leave=False) # 进度条设置
    for images, labels in pbar:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad() # 清空之前残留的梯度值，防止梯度再次累加。
        outputs = model(images) # 正向传播
        loss = criterion(outputs, labels) # 计算损失
        loss.backward() # 计算当前梯度（反向传播）
        # 在 PyTorch 中，每次 .backward() 计算梯度时，默认是将当前梯度加到已有梯度上，而不是替换。
        optimizer.step() # 用梯度更新参数

        # 统计
        running_loss += loss.item() * images.size(0) # 当前 epoch 累计损失总和
        # loss.item()：将当前 batch 的 loss 从 Tensor 转成 Python float（否则不能累加）。images.size(0)：当前 batch 中图像的数量（也就是 batch_size）。为什么乘以 batch size？因为 loss 是对当前 batch 的“平均损失”，我们要恢复成“总损失”（总和），方便最后平均
        _, predicted = torch.max(outputs, 1)
        # 取每张图预测结果中概率最大的类别作为最终预测。outputs 是模型输出的 logits（形状：[batch_size, num_classes]）。返回 (最大值, 索引)，我们只关心索引，即类别编号。predicted 是形状为 [batch_size] 的预测标签列表
        
        correct += (predicted == labels).sum().item()
        # 统计当前 batch 中预测正确的样本个数。(predicted == labels) 得到一个布尔张量，比如 [True, False, True]。.sum() 对布尔张量求和，相当于统计正确预测个数。.item() 把张量转成 Python 数字，便于加到 correct 变量里
        
        total += labels.size(0) # 统计这个 batch 的样本数，加到总数中

        # 更新进度条
        pbar.set_postfix(loss=loss.item())

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

# 验证函数
def validate(model, loader, criterion):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    pbar = tqdm(loader, desc='Validating', leave=False)
    with torch.no_grad(): # 在不需要反向传播的阶段（如验证、测试、推理），关闭 PyTorch 的自动求导机制，节省内存和计算。
    # 在训练过程中：PyTorch 会自动构建计算图（computation graph）并保存所有中间变量（用来 .backward() 反向传播）
    # 在验证或测试时：我们只是前向传播（forward）看模型效果，不需要计算梯度，也不会调用 .backward()。
        for images, labels in pbar:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            correct += (predicted == labels).sum().item()
            total += labels.size(0)

            pbar.set_postfix(loss=loss.item())

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

# 训练循环
best_val_acc = 0.0
for epoch in range(EPOCHS):
    print(f"\nEpoch {epoch+1}/{EPOCHS}")

    train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion)
    val_loss, val_acc = validate(model, test_loader, criterion)

    print(f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
          f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")

    # 保存最佳模型
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), SAVE_PATH)
        print(f"保存新模型：{SAVE_PATH}，验证准确率提升为 {val_acc:.4f}")

print("训练完成！最佳验证准确率为：", best_val_acc)

import os
import torch
import random
import matplotlib.pyplot as plt
import matplotlib
from torch.utils.data import DataLoader
from dataset import CatDogDataset
from model import SimpleCNN
import torchvision.transforms as transforms
from sklearn.metrics import classification_report

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 参数
BATCH_SIZE = 32
IMG_SIZE = 224
MODEL_PATH = 'best_model.pth'
NUM_VIS = 8  # 可视化图像数

# 类别标签映射
idx_to_class = {0: 'cat', 1: 'dog'}

# 测试集预处理
test_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])

# 加载测试集
test_dataset = CatDogDataset(root_dir='data/test', transform=test_transform)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# 加载模型
model = SimpleCNN(num_classes=2).to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()

# 评估准确率
correct = 0
total = 0
# classification report
all_preds = []
all_labels = []

with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        correct += (predicted == labels).sum().item()
        total += labels.size(0)
        
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

acc = correct / total
print(f"Test Accuracy: {acc * 100:.2f}%")

# 打印分类报告
print("\n Classification Report:")
print(classification_report(all_labels, all_preds, target_names=['cat', 'dog']))

# 可视化部分预测结果
# 反归一化函数（将[-1, 1]映射回[0, 1]）
def denormalize(tensor):
    return tensor * 0.5 + 0.5

# 随机挑选样本
sample_indices = random.sample(range(len(test_dataset)), NUM_VIS)

plt.figure(figsize=(16, 8))
for i, idx in enumerate(sample_indices):
    image, label = test_dataset[idx]
    input_tensor = image.unsqueeze(0).to(device) # 模型输入 shape = [batch_size, 3, 224, 224]，所以在维度0的位置增加一个维度，image.shape = [3, 224, 224] 变成 [1, 3, 224, 224]

    with torch.no_grad():
        output = model(input_tensor)
        _, pred = torch.max(output, 1)

    # 转换为可显示格式
    image_np = denormalize(image).permute(1, 2, 0).cpu().numpy()
    # .permute(1, 2, 0)，PyTorch 的图像张量默认顺序是：(C, H, W) = (通道, 高, 宽)。而 Matplotlib/NumPy 的图像格式是：(H, W, C)，所以 .permute(1, 2, 0) 就是交换维度。
    # matplotlib 没法直接访问，需要先转到 CPU

    plt.subplot(2, NUM_VIS // 2, i + 1)
    plt.imshow(image_np)
    plt.axis('off')
    plt.title(f"Pred: {idx_to_class[pred.item()]}\nTrue: {idx_to_class[label]}", fontsize=12)

plt.tight_layout()
plt.savefig("prediction_visualization.png")
plt.show()

import torch
import torch.nn as nn
import torchvision.models as models

class FineTunedResNet18(nn.Module):
    """
    使用ImageNet预训练的ResNet18模型进行微调
    - 冻结所有卷积层
    - 替换最后的分类层为2分类
    """
    def __init__(self, num_classes=2, freeze_features=True):
        super(FineTunedResNet18, self).__init__()

        # 加载预训练的 ResNet18
        self.model = models.resnet18(pretrained=True)

        # 是否冻结特征提取部分
        if freeze_features:
            for param in self.model.parameters(): #返回模型中所有的可学习参数（通常是权重 weight 和偏置 bias）。也就是：ResNet18 的卷积层、BN层、全连接层等的参数列表。
                param.requires_grad = False

        # 替换最后的全连接层
        in_features = self.model.fc.in_features
        # 获取原始 fc 层的输入维度（通常为 512）；这个值不变，我们用它来构造一个新的输出层。
        self.model.fc = nn.Linear(in_features, num_classes)

    def forward(self, x):
        return self.model(x)

# from model import SimpleCNN
from model import FineTunedResNet18
# 初始化微调模型（默认冻结特征层）
model = FineTunedResNet18(num_classes=2, freeze_features=True).to(device)  

from model import FineTunedResNet18
model = FineTunedResNet18(num_classes=2, freeze_features=True).to(device)