#深度学习视觉全通关：从CNN原理到产品化部署

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

class BasicCNN(nn.Module):
    """
    基础CNN网络结构示例
    """
    def __init__(self):
        super(BasicCNN, self).__init__()
        
        # 卷积层 1: 提取低级特征（边缘、纹理）
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        
        # 卷积层 2: 提取中级特征
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        
        # 最大池化层
        self.pool = nn.MaxPool2d(2, 2)
        
        # 全连接层
        self.fc1 = nn.Linear(64 * 56 * 56, 128)  # 假设输入为224x224
        self.fc2 = nn.Linear(128, 2)  # 二分类
        
        # Dropout防止过拟合
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        # 第一个卷积块
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        # 第二个卷积块
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        
        # 展平
        x = x.view(-1, 64 * 56 * 56)
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 创建模型实例
model = BasicCNN()
print("基础CNN模型创建完成")

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import os

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

# 定义数据增强
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet标准化
])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 假设数据集在'data/dogs_vs_cats'目录下，结构如下：
# data/dogs_vs_cats/
#   ├── train/
#   │   ├── cats/
#   │   └── dogs/
#   └── val/
#       ├── cats/
#       └── dogs/

# 加载数据集
try:
    train_data = datasets.ImageFolder(root='data/dogs_vs_cats/train', transform=train_transform)
    val_data = datasets.ImageFolder(root='data/dogs_vs_cats/val', transform=val_transform)
    
    train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_data, batch_size=32, shuffle=False, num_workers=4)
    
    print(f"训练集大小: {len(train_data)}")
    print(f"验证集大小: {len(val_data)}")
    print(f"类别数量: {len(train_data.classes)}")
    print(f"类别名称: {train_data.classes}")
except FileNotFoundError:
    print("数据集未找到，请确保数据集路径正确")
    # 创建模拟数据集用于演示
    print("使用模拟数据进行演示...")

class CustomCNN(nn.Module):
    """
    自定义CNN模型
    """
    def __init__(self, num_classes=2):
        super(CustomCNN, self).__init__()
        
        # 第一个卷积块
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        
        # 第二个卷积块
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        
        # 第三个卷积块
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        
        # 最大池化层
        self.pool = nn.MaxPool2d(2, 2)
        
        # Dropout层
        self.dropout = nn.Dropout(0.5)
        
        # 全连接层
        # 经过3次池化后，224x224 -> 28x28
        self.fc1 = nn.Linear(128 * 28 * 28, 512)
        self.fc2 = nn.Linear(512, 128)
        self.fc3 = nn.Linear(128, num_classes)
        
    def forward(self, x):
        # 第一个卷积块
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        
        # 第二个卷积块
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        
        # 第三个卷积块
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        
        # 展平
        x = x.view(-1, 128 * 28 * 28)
        
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        
        return x

# 创建模型实例
custom_model = CustomCNN(num_classes=2).to(device)
print("自定义CNN模型创建完成")
print(f"模型参数数量: {sum(p.numel() for p in custom_model.parameters()):,}")

def train_model(model, train_loader, val_loader, epochs=10, learning_rate=0.001):
    """
    训练模型
    """
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
    
    best_val_acc = 0.0
    
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        running_loss = 0.0
        correct_train = 0
        total_train = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # 统计
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total_train += labels.size(0)
            correct_train += (predicted == labels).sum().item()
        
        # 验证阶段
        model.eval()
        correct_val = 0
        total_val = 0
        val_loss = 0.0
        
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()
                
                _, predicted = torch.max(outputs, 1)
                total_val += labels.size(0)
                correct_val += (predicted == labels).sum().item()
        
        # 计算准确率
        train_acc = 100 * correct_train / total_train
        val_acc = 100 * correct_val / total_val
        avg_train_loss = running_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)
        
        print(f'Epoch [{epoch+1}/{epochs}]')
        print(f'Train Loss: {avg_train_loss:.4f}, Train Acc: {train_acc:.2f}%')
        print(f'Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%')
        print('-' * 50)
        
        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_custom_cnn.pth')
            print(f'--> 检测到更好的模型，已保存 (Val Acc: {best_val_acc:.2f}%)')
        
        # 更新学习率
        scheduler.step()
    
    print(f'训练完成！最佳验证准确率: {best_val_acc:.2f}%')
    return best_val_acc

# 如果有数据，可以运行训练
# best_acc = train_model(custom_model, train_loader, val_loader, epochs=10)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import os

def predict_local_image(image_path, model_path='best_custom_cnn.pth', num_classes=2):
    """
    预测本地图片
    """
    # 检查模型文件是否存在
    if not os.path.exists(model_path):
        print(f"错误：找不到模型文件 {model_path}")
        return
    
    # 准备设备和模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CustomCNN(num_classes=num_classes).to(device)
    
    # 加载训练好的权重
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()  # 切换到评估模式（关闭 Dropout）
        print("模型权重加载成功！")
    except Exception as e:
        print(f"模型加载失败: {e}")
        return

    # 定义与训练一致的预处理
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 加载并转换图片
    try:
        img = Image.open(image_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(device)  # 增加 batch 维度 [1, 3, 224, 224]
    except Exception as e:
        print(f"图片读取失败: {e}")
        return

    # 执行推理
    with torch.no_grad():  # 推理不需要计算梯度
        outputs = model(img_tensor)
        probabilities = F.softmax(outputs, dim=1)  # 转化为概率
        confidence, predicted = torch.max(probabilities, 1)

    # 输出结果
    classes = ['猫 (Cat)', '狗 (Dog)']
    result = classes[predicted.item()]
    confidence_percent = confidence.item() * 100
    
    print("=" * 40)
    print(f"预测结果: {result}")
    print(f"置信度: {confidence_percent:.2f}%")
    print(f"各类别概率:")
    for i, class_name in enumerate(classes):
        prob = probabilities[0][i].item() * 100
        print(f"  {class_name}: {prob:.2f}%")
    print("=" * 40)

# 使用示例（取消注释以运行）
# predict_local_image('path/to/your/image.jpg')

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader

def create_resnet_model(num_classes=2):
    """
    创建基于ResNet18的迁移学习模型
    """
    # 加载预训练的ResNet18模型
    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    
    # 冻结预训练层（可选）
    # for param in model.parameters():
    #     param.requires_grad = False
    
    # 修改最后的全连接层
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)
    
    return model

def train_resnet_transfer(epochs=10):
    """
    训练ResNet迁移学习模型
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 数据预处理（与之前相同）
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(degrees=15),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    # 创建模型
    model = create_resnet_model(num_classes=2).to(device)
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    # 使用较小的学习率，因为预训练模型已经很好了
    optimizer = optim.Adam(model.parameters(), lr=0.0001, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
    
    # 训练循环（简化版，实际使用时需要加载数据）
    print("ResNet18迁移学习模型已准备就绪")
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    print(f"可训练参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
    
    # 保存模型结构
    torch.save(model.state_dict(), 'resnet18_transfer_init.pth')
    
    return model

# 创建ResNet模型
resnet_model = train_resnet_transfer()

def predict_with_resnet(image_path, model_path='resnet18_transfer_init.pth'):
    """
    使用ResNet模型预测图片
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 创建模型结构
    model = models.resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, 2)  # 2分类

    # 加载权重
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model = model.to(device)
        model.eval()
        print("ResNet模型权重加载成功！")
    except Exception as e:
        print(f"ResNet模型加载失败: {e}")
        return

    # 预处理
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 读取与预测
    try:
        img = Image.open(image_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(device)

        with torch.no_grad():
            outputs = model(img_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)

        classes = ['猫 (Cat)', '狗 (Dog)']
        result = classes[predicted.item()]
        confidence_percent = confidence.item() * 100
        
        print(f"ResNet18预测结果: {result}")
        print(f"置信度: {confidence_percent:.2f}%")
        
        return result, confidence_percent
    except Exception as e:
        print(f"预测失败: {e}")
        return None, 0

# 使用示例（取消注释以运行）
# predict_with_resnet('path/to/your/image.jpg')

from torchvision import models

def create_vit_model(num_classes=2):
    """
    创建基于Vision Transformer的模型
    """
    # 加载预训练的ViT Base模型
    weights = models.ViT_B_16_Weights.DEFAULT
    model = models.vit_b_16(weights=weights)
    
    # 修改分类头
    num_in_features = model.heads.head.in_features
    model.heads.head = nn.Linear(num_in_features, num_classes)
    
    return model

def train_vit_model(epochs=5):
    """
    训练ViT模型
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # ViT专用预处理
    vit_transform = transforms.Compose([
        transforms.Resize((224, 224)),  # ViT通常使用224x224
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # ViT常用归一化
    ])
    
    # 创建ViT模型
    model = create_vit_model(num_classes=2).to(device)
    
    # 定义损失函数和优化器（ViT通常使用更小的学习率）
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-5, weight_decay=0.1)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    print("Vision Transformer模型已准备就绪")
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # 保存初始模型
    torch.save(model.state_dict(), 'vit_initial.pth')
    
    return model

# 创建ViT模型
vit_model = train_vit_model()

def predict_with_vit(image_path, model_path='vit_initial.pth'):
    """
    使用ViT模型预测图片
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 创建ViT模型结构
    model = models.vit_b_16(weights=None)
    model.heads.head = nn.Linear(model.heads.head.in_features, 2)

    # 加载权重
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model = model.to(device)
        model.eval()
        print("ViT模型权重加载成功！")
    except Exception as e:
        print(f"ViT模型加载失败: {e}")
        return

    # ViT专用预处理
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

    # 预测
    try:
        img = Image.open(image_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(device)

        with torch.no_grad():
            output = model(img_tensor)
            prob = F.softmax(output, dim=1)
            score, pred = torch.max(prob, 1)

        classes = ['猫 (Cat)', '狗 (Dog)']
        result = classes[pred.item()]
        confidence = score.item() * 100
        
        print(f"ViT预测: {result} (置信度: {confidence:.2f}%)")
        
        return result, confidence
    except Exception as e:
        print(f"ViT预测失败: {e}")
        return None, 0

# 使用示例（取消注释以运行）
# predict_with_vit('path/to/your/image.jpg')

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
import gradio as gr
from PIL import Image
import numpy as np

def load_model(model_type='resnet', model_path=None):
    """
    加载训练好的模型
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    if model_type == 'resnet':
        model = models.resnet18(weights=None)
        model.fc = nn.Linear(model.fc.in_features, 2)
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    elif model_type == 'vit':
        model = models.vit_b_16(weights=None)
        model.heads.head = nn.Linear(model.heads.head.in_features, 2)
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])
    else:  # custom cnn
        model = CustomCNN(num_classes=2)
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    
    # 加载权重
    if model_path:
        model.load_state_dict(torch.load(model_path, map_location=device))
    
    model.to(device).eval()
    return model, transform, device

def classify_image(img, model_type='resnet'):
    """
    分类函数
    """
    if img is None:
        return {"错误": 1.0}
    
    # 加载模型
    model, transform, device = load_model(model_type)
    
    # 预处理
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype('uint8'), 'RGB')
    
    img_tensor = transform(img).unsqueeze(0).to(device)
    
    # 推理
    with torch.no_grad():
        output = model(img_tensor)
        probabilities = F.softmax(output, dim=1)[0]
    
    # 返回结果
    classes = ['猫 (Cat)', '狗 (Dog)']
    results = {}
    for i, class_name in enumerate(classes):
        results[class_name] = float(probabilities[i])
    
    return results

# 创建Gradio界面
with gr.Blocks(title="智能猫狗识别系统") as demo:
    gr.Markdown("# 🐱 智能猫狗识别系统 🐶")
    gr.Markdown("上传一张猫或狗的照片，AI会告诉你它的类别及置信度。")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="上传图片", type="numpy")
            model_choice = gr.Radio(
                choices=["CustomCNN", "ResNet18", "Vision Transformer"],
                value="ResNet18",
                label="选择模型"
            )
            run_button = gr.Button("识别")
        
        with gr.Column():
            output_label = gr.Label(label="识别结果")
    
    run_button.click(
        fn=classify_image,
        inputs=[image_input, model_choice],
        outputs=output_label
    )

# 启动应用
def launch_app():
    """
    启动Gradio应用
    """
    print("启动猫狗识别Web应用...")
    print("本地访问地址: http://localhost:7860")
    print("设置share=True可生成公网访问链接")
    
    # 在实际使用时取消注释下面这行
    # demo.launch(share=True)

# launch_app()  # 取消注释以启动应用