#Vision Transformer：从图像切片到Patch Embedding详解

import torch
import torch.nn as nn
from einops import rearrange

class ImageToPatches(nn.Module):
    """
    图像到 patch 的转换模块
    """
    def __init__(self, image_size=224, patch_size=16, channels=3):
        super().__init__()
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        self.patch_dim = channels * patch_size ** 2

        # 线性投影层：将每个 patch 的原始像素映射到目标维度
        self.projection = nn.Linear(self.patch_dim, self.patch_dim)

    def forward(self, x):
        """
        x: (batch, channels, height, width)
        return: (batch, num_patches, patch_dim)
        """
        batch_size, channels, height, width = x.shape

        # 简单校验输入尺寸
        assert height == self.image_size and width == self.image_size, \
            f"输入图像尺寸应为 ({self.image_size}, {self.image_size})"

        # 使用 einops 的 rearrange 优雅切分
        x = rearrange(
            x,
            'b c (h p1) (w p2) -> b (h w) (p1 p2 c)',
            p1=self.patch_size,
            p2=self.patch_size
        )

        # 线性投影
        x = self.projection(x)

        return x

import torch
import torch.nn as nn

class VisionTransformer(nn.Module):
    """
    Vision Transformer 完整实现
    """
    def __init__(
        self,
        image_size=224,
        patch_size=16,
        num_classes=1000,
        dim=768,             # 嵌入维度
        depth=12,            # Transformer 层数
        heads=12,            # 注意力头数
        mlp_dim=3072,        # 前馈网络扩展维度
        dropout=0.1,
        emb_dropout=0.1
    ):
        super(VisionTransformer, self).__init__()
        num_patches = (image_size // patch_size) ** 2
        patch_dim = 3 * patch_size ** 2

        # 用卷积一步完成切块+投影 (卷积核大小和步长都等于 patch_size)
        self.to_patch_embedding = nn.Sequential(
            nn.Conv2d(3, patch_dim, kernel_size=patch_size, stride=patch_size),
            nn.Flatten(start_dim=2),            # (B, patch_dim, H, W) -> (B, patch_dim, num_patches)
            nn.Linear(patch_dim, dim)
        )

        # 可学习的分类 token 和位置编码
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

        self.dropout = nn.Dropout(emb_dropout)

        # Transformer 编码器（使用 PyTorch 官方的 TransformerEncoderLayer）
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=dim,
            nhead=heads,
            dim_feedforward=mlp_dim,
            dropout=dropout,
            activation='gelu',
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=depth)

        # 最后的分类头
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )

    def forward(self, img):
        # 1. 获取 patch 嵌入
        x = self.to_patch_embedding(img)    # (B, num_patches, dim)
        b, n, _ = x.shape

        # 2. 在最前面拼接 class token
        cls_tokens = self.cls_token.repeat(b, 1, 1)   # (B, 1, dim)
        x = torch.cat([cls_tokens, x], dim=1)         # (B, num_patches+1, dim)

        # 3. 加上位置编码
        x = x + self.pos_embedding[:, :(n + 1)]
        x = self.dropout(x)

        # 4. 送入 Transformer 编码器
        x = self.transformer(x)

        # 5. 取出 class token 对应位置的输出进行分类
        cls_output = x[:, 0]    # (B, dim)
        output = self.mlp_head(cls_output)

        return output

import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    """
    多头自注意力机制
    """
    def __init__(self, d_model=768, num_heads=12, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 生成 Q, K, V 的线性变换层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)
        # 缩放因子
        self.scale = torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # 线性投影
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # 拆分成多头: (B, seq_len, num_heads, d_k) -> (B, num_heads, seq_len, d_k)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        attention = F.softmax(scores, dim=-1)
        attention = self.dropout(attention)

        # 加权求和
        output = torch.matmul(attention, V)   # (B, num_heads, seq_len, d_k)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

        # 最终线性变换
        output = self.W_o(output)

        return output

import torch
from torchvision import models, transforms

# 加载预训练 ViT-B/16
model = models.vit_b_16(weights='IMAGENET1K_V1')
model.eval()

# 图像预处理
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

# 假设已经得到 input_tensor，Shape = (1, 3, 224, 224)
with torch.no_grad():
    output = model(input_tensor)
    probabilities = torch.nn.functional.softmax(output[0], dim=0)

from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests

# 加载处理器和模型
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

# 读取图片
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)

inputs = processor(images=image, return_tensors="pt")

# 推理
outputs = model(**inputs)
predicted_class_idx = outputs.logits.argmax(-1).item()
print("预测类别:", model.config.id2label[predicted_class_idx])